Paper: MOTOR: A Time-to-Event Foundation Model For Structured Medical Records (ICLR-2024)

Reviewer comment (Novelty):
If using the fact that ICD code sequences are sequences of events rather than observation to directly perform large-scale multitask dynamic survival analysis on them is great for this type of data, it is also very limiting to this unique type of data. Indeed, as mentioned by the authors, survival analysis is useful in a variety of domains ranging from cancer research to finance. However, in such domains, to my knowledge, data cannot be formalized as a sequence of events, hence the proposed pretraining task could not be expanded to other time-to-event models.

Author response:
> The specificity of the task to ICD codes data

Our method is applicable to datasets that have timestamped event information. Such datasets are not restricted to medicine and do exist in enterprise settings. For example, user activity datasets that track user interactions within a particular company / software app (such as [1] and [2]) have similar streams of timestamped event data that may be suitable for our method.  One challenge is that these timestamped, event datasets are difficult for researchers to obtain at scale because companies generally view them as high-value, proprietary data. 

Medical records offer the advantage of being available at the scale necessary for training foundation models and are generally more accessible to researchers (e.g., MIMIC-IV, MERATIVE). However, we are hopeful that future work will be able to reproduce our results on other sources of timestamped event datasets.


1. How Jing and Alexander J. Smola. 2017. Neural Survival Recommender. In Proceedings of the Tenth ACM International Conference on Web Search and Data Mining (WSDM '17). Association for Computing Machinery, New York, NY, USA, 515–524. https://doi.org/10.1145/3018661.3018719

2. Á. Periáñez, A. Saas, A. Guitart and C. Magne, "Churn Prediction in Mobile Social Games: Towards a Complete Assessment Using Survival Ensembles," 2016 IEEE International Conference on Data Science and Advanced Analytics (DSAA), Montreal, QC, Canada, 2016, pp. 564-573, doi: 10.1109/DSAA.2016.84.

Reviewer comment (Experiments):
The authors pre-train their model in the so-called "dynamic survival analysis" setting, however, they only evaluate it in a static setting. It would have been a great addition to also have dynamic tasks. The same baselines can be considered with a landmarking approach.

Author response:
> The lack of experiments in the dynamic setting


We agree that more exploration of dynamic vs static survival analysis settings is an important research direction and that the lack of dynamic setting evaluations is a limitation. We hope to explore these types of evaluation settings in future work.

Reviewer comment (Evaluation):
The attention of the authors to ethical concerns is clearly above standard. As part of it, they perform a sub-group analysis per ethnicity. Compared to the RSF model which is quite stable across groups, their model exhibits a significant drop in performance among Native Americans (~5% lower than the closest group). However, the authors have the following statement: "We find that with one statistically insignificant exception, MOTOR-Finetune does not reduce the performance within sensitive groups". I have a hard time believing that such a drop is "statistically insignificant". I believe the authors should clearly state this limitation.

Author response:
> The drop in performance for the Native American sub-group

We apologize for not providing sufficient details to make a convincing case. The issue is that the Native American subgroup is very small in EHR-OMOP, only 0.3% of the total patients fall into that category, so our estimates for that group are very wide. We have explicitly added confidence intervals to Appendix J.12 (table reproduced below) to better illustrate the issue. Note how the confidence interval for the Native American sub-group contains 0, and is thus not statistically significant.

The table below contains the average time-dependent C statistic for RSF, MOTOR-Finetune, and the difference between the two models within important subgroups. The 95\% confidence interval for the difference is also included. Bolding indicates the best-performing model. $\ast$ indicates statistically significant entries at p = 0.05.

|Group|RSF|MOTOR-Finetune|Delta [Confidence Interval]|
|-|-|-|-|
|Female|0.799|**0.856**|0.058 [0.049, 0.067]\*|
|Male|0.790|**0.845**|0.055 [0.041, 0.069]\*|
|-|-|-|-|
|American Indian or Alaska Native|**0.810**|0.787|-0.022 [-0.209, 0.171]|
|Asian|0.793|**0.853**|0.059 [0.039, 0.081]\*|
|African American|0.761|**0.834**|0.073 [0.006, 0.132]\*|
|Native Hawaiian or Other Pacific Islander|0.841|**0.914**|0.073 [-0.044, 0.154]|
|White|0.788|**0.850**|0.062 [0.054, 0.070]\*|
|Unknown Race|0.832|**0.887**|0.055 [0.039, 0.073]\*|
|-|-|-|-|
|Unknown Ethnicity|0.847|**0.883**|0.036 [0.014, 0.065]\*|
|Hispanic|0.794|**0.857**|0.063 [0.041, 0.087]\*|
|Non-Hispanic|0.794|**0.853**|0.059 [0.052, 0.066]\*|



The differences for the Native Hawaiin or Other Pacific Islander racial category are also not statistically significant, which reflects how only 1% of patients in EHR-OMOP belong to that racial group.

We agree that performance in these groups is important, so we have added an explicit disclaimer in our model card that our model is not tested well on those two subgroups. We hope that our model release will help address this problem by allowing other researchers with access to those populations to validate our model for those groups.

Reviewer comment (Reproducibility):
How does the author handle events that can occur multiple times? Do they predict the time of the next event instead?

Author response:
> How does the author handle events that can occur multiple times? Do they predict the time of the next event instead?

We predict the time until the next event for events that can occur multiple times. We have modified the relevant sentence in the methods section to make this more clear.

Reviewer comment (Experiments):
Limited ablation on the amount of pretraining data (e.g., 1%, 10%). Although the authors have made a significant contribution by releasing the model, providing more details on the amount of data necessary to achieve the current performance would be highly beneficial. Time-to-event modeling holds importance not only in medical settings but also in various other domains.

Author response:
We agree that a deeper analysis of TTE pretraining, including more detailed analyses of pretraining sizes, is an exciting direction for further exploration. In the presence of limited compute budgets, our current pretraining sweep shown in Figure 4 used 5%, 10%, 25%, and 100% of the input data. Our pretraining ablations are able to highlight 2 points: 

(1) In EHR data (EHR-OMOP), we observed that pre-training provides larger gains when using more data, suggesting potentially further gains beyond the 2.5M patients that comprised our entire EHR patient pre-training set.

(2) In claims data (MERATIVE), we observed some benefits from pre-training beyond our smallest sample (2.74M patients), but overall total benefit was 89% less than the gain from EHR data.

These findings provide some guidance on how much data practitioners need in that small EHR datasets are insufficient to pretrain MOTOR effectively and pretraining MOTOR on very large claims datasets does not provide significant additional predictive value.

Reviewer comment (Presentation):
Figure 1 could be explained or illustrated in more detail. For example, it is hard to understand what each survival curve means for each timestep in pretraining tasks.

Author response:
Thank you for pointing out the need for Figure 1 needing more explanation and better element labels. For the survival plots, where the x axis is time and the y axis is the probability, each time step adds new information that can update the estimate of y in the survival curve. 

We have updated the figure capture to explain this intuition and better label figure elements.

Reviewer comment (Reproducibility):
How much time did it take to pretrain the model?

Author response:
Pretraining on EHR-OMOP takes 111 V100 GPU hours. Pretraining on MERATIVE takes 330 V100 GPU hours. We have added this detail to the compute section in Appendix C.

Reviewer comment (Experiments):
What is the source of pretraining data?

Author response:
We pretrain on two datasets, EHR-OMOP and MERATIVE. EHR-OMOP is a private dataset provided by an academic medical center. MERATIVE is commercially available and our institution purchased a license.

Reviewer comment (Experiments):
Can you elaborate more on the 8,192 code selection process?

Author response:
We have added Appendix O with a more comprehensive description of our code selection process. We have reproduced a summarized version of that new section here. The core idea is to do code selection using a simple heuristic that, as closely as possible, matches a standard frequency-based selection technique (i.e., including the top-k most common codes for pretraining) but accounts for the ontology relationships between codes. In the absence of ontology relationships, ranking using frequency is equivalent to ranking by Shannon entropy. Therefore, to have a ranking criteria that is similar to frequency-based selection, but taking into account ontologies, we use the conditional Shannon entropy formula, conditioning on the presence or absence of the parents of a particular code.

Let $C$ be a boolean random variable indicating the presence or absence of a particular code in a patient timeline and let $O$ be a boolean random variable indicating the presence or absence of the parents of that code.

The conditional entropy formula is:

$$H(C|O)\ = -\sum_{c\in\mathcal C, o\in\mathcal O}p(o,c)\log \frac {p(o, c)} {p(o)}$$

This simplifies to:

\begin{equation}
H(C|O) =    -p(O = T,C = F)\log \frac{p(O = T, C = F)} {p(O = T)} +  -p(O = T,C = T)\log \frac{p(O = T, C = T)} {p(O = T)}
\end{equation}

This formula can then be evaluated by plugging in $p(O = T, C = F)$, $p(O = T, C = T)$, and $p(O = T)$, all of which can be estimated by counting the frequency of codes (conditioned on the presence / absence of their parents) in the training data. The highest entropy codes according to this formula are then selected.

Reviewer comment (Experiments):
Is the patient data at the admission level or concatenated across admissions?

Author response:
The patient data is concatenated across admissions.

Reviewer comment (Experiments):
It is not clear how the authors handle missing information in their analyses. How are missing values in the time-series data represented?

Author response:
We acknowledge that missing data (particularly informative missingness) is an important aspect of modeling EHR data. As noted in Appendix F, we follow BEHRT[1]’s general Transformer approach, which attempts to handle missingness using a sequence formulation, where patient information is received one event at a time, events with missing data are not generated. Because event times are provided as a feature to a deep neural network architecture that can process complex feature interactions, in principle the model is capable of learning informative patterns of missingness when certain events are systematically absent. 

We have updated Appendix F to include the above information.

1. Li, Y., Rao, S., Solares, J.R.A. et al. BEHRT: Transformer for Electronic Health Records. Sci Rep 10, 7155 (2020). https://doi.org/10.1038/s41598-020-62922-y

Reviewer comment (Reproducibility):
While the authors explain the subsampling process in the Appendix, this is critical and forms the basis of the pre-training approach. Expanding on the rationale/method for this seems important.

Author response:
Thank you for the suggestion to expand on this matter. The subsampling approach noted in Section 4 is not used in any way as part of our pre-training process. Subsampling is only performed when creating target task datasets for evaluation purposes, given the scale of our evaluation datasets (i.e., containing up to tens of millions of patients), which makes it infeasible to evaluate more computationally expensive baseline methods like RSF on the entire dataset. Section 4 and Appendix D elaborate on this motivation. 
While this subsampling does not impact the relative ranking and comparisons of methods in our manuscript, it does induce a non-uniform scaling of the hazard rates. Since calculating the unbiased hazard rate is often useful, Appendix D also formally outlines the exact hazard rate scaling induced by our evaluation subsampling.

Reviewer comment (Reproducibility):
How are missing values in the time-series data represented?

Author response:
We acknowledge that missing data (particularly informative missingness) is an important aspect of modeling EHR data. As noted in Appendix F, we follow BEHRT[1]’s general Transformer approach, which attempts to handle missingness using a sequence formulation, where patient information is received one event at a time, events with missing data are not generated. Because event times are provided as a feature to a deep neural network architecture that can process complex feature interactions, in principle the model is capable of learning informative patterns of missingness when certain events are systematically absent. 

We have updated Appendix F to include the above information.

1. Li, Y., Rao, S., Solares, J.R.A. et al. BEHRT: Transformer for Electronic Health Records. Sci Rep 10, 7155 (2020). https://doi.org/10.1038/s41598-020-62922-y

Reviewer comment (Reproducibility):
Are specific transformer base models used?

Author response:
Sorry for the confusion. We specified this in the appendix, but not in the main body of the paper. We use a local attention transformer base model, with the exact implementation details matching the local attention base model described in [1]. We have since moved this detail and reference into the main body of the paper.

1. Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53–68, 2021. doi: 10.1162/tacl_a_00353. URL https://aclanthology.org/2021.tacl-1.4

Reviewer comment (Presentation):
What are the x and y axis in Figure 1 - pretraining tasks? Are they hazard ratio curves?

Author response:
We acknowledge this figure was confusing in our original submission. These are survival curves, with the x axis as time and the y axis as the probability that the event occurs after x time. We have updated the figure to more clearly label the axes and other figure elements.

Reviewer comment (Evaluation):
How is the time-to-event task evaluated on the MIMIC dataset? Is the goal to predict the diagnosis code at each ICU admission?

Author response:
We take advantage of the longitudinal non-ICU data provided by MIMIC-IV, which allows us to use the same task setup as within MERATIVE and EHR-OMOP, aka predicting the time until the first diagnosis of a condition from the end of a given visit. We have updated Appendix J.7 (which describes the MIMIC experiments) to make this more clear.

Reviewer comment (Reproducibility):
How many resources are needed to fine-tune or infer using the pretrained model?

Author response:
Our timing experiments (which are detailed in Appendix J.10) show that our model has a throughput of 147 patients per second on EHR-OMOP when performing inference using a single NVIDIA V100. We have added that throughput number to Appendix J.10.