Adaptive Discretization for Event PredicTion (ADEPT)

Abstract

Recently developed survival analysis methods improve upon existing approaches by predicting the probability of event occurrence in each of a number pre-specified (discrete) time intervals. By avoiding placing strong parametric assumptions on the event density, this approach tends to improve prediction performance, particularly when data are plentiful. However, in clinical settings with limited available data, it is often preferable to judiciously partition the event time space into a limited number of intervals well suited to the prediction task at hand. In this work, we develop Adaptive Discretization for Event PredicTion (ADEPT) to learn from data a set of cut points defining such a partition. We show that in two simulated datasets, we are able to recover intervals that match the underlying generative model. We then demonstrate improved prediction performance on three real-world observational datasets, including a large, newly harmonized stroke risk prediction dataset. Finally, we argue that our approach facilitates clinical decision-making by suggesting time intervals that are most appropriate for each task, in the sense that they facilitate more accurate risk prediction.

1. INTRODUCTION

Time to event modeling, also called survival analysis, is ubiquitous throughout clinical medicine as well as in many other fields concerned with predicting risk of events of interest (e.g., clinical outcomes) based on available features (e.g., patient characteristics). Traditional approaches include the well-known Cox proportional hazards (Cox-PH) model [Cox, 1972], in which features modulate a baseline hazard rate; and the accelerated failure time (AFT) model [Wei, 1992] model, in which features accelerate or decelerate a learned, parametric event time density.

Recently developed methods have focused on a) allowing effects of features on the hazard rate or event time density to be non-linear and flexible [Katzman et al., 2018, Ranganath et al., 2016, Kvamme et al., 2019, Miscouridou et al., 2018]; and b) also allowing greater flexibility in the form of the event time density itself via approaches that discretize time, then predict the probability of event occurrence in each resulting time interval [Yu et al., 2011, Lee et al., 2018, Ren et al., 2019, Tjandra et al., 2021, Engelhard and Henao, 2022].

The prognostic information provided by these models often has direct and significant impact on stakeholder decision-making. In a clinical setting, for example, information about risk within a particular time interval might influence providers’ or patients’ decisions about whether to pursue treatment, or which specific treatment to pursue. It is therefore critical not only that predictions are accurate, but also that they are easily interpretable by stakeholders who wish to integrate them in decision-making. The predictions of a Cox-PH model might be presented to stakeholders as relative hazards, for instance, whereas it is natural to present the predictions of more recent models as the probability of event occurrence in a time interval of interest.

Importantly, however, decisions about these intervals made during model development – in other words, choices about the number and placement of cut points used to discretize the event time space – can have substantial impact on interpretability as well as performance. Equipped with unlimited data, we might use a large number of cut points to divide the timeline into tiny intervals; this would then allow us to summarize risk over an arbitrary time period of interest by combining predictions across all the intervals that comprise that period. However, the amount of data required to accurately estimate risk in each interval increases as the number of intervals increases, making this approach impractical even for large observational datasets. Equipped with unlimited time, on the other hand, we might present risk in a format most relevant to a particular patient, or to the decision at hand. Again, however, practical considerations typically require us to instead summarize risk over a consistent, limited number of time-frames (e.g., 10-year risk, 5-year risk). In some cases a particular discretization is most actionable given the clinical context, but in others the choice is arbitrary, and it would be preferable to identify a discretization that facilitates more accurate prediction.

To illustrate the problem more concretely, consider the following example from the maternal health setting, which partly motivated this work. Patients with preeclampsia and gestational hypertension have substantially increased risk of postpartum cardiovascular events [Meng et al., 2022], but this risk can be mitigated by regular monitoring (e.g., increased visits) of high-risk patients in the months after delivery. When developing a monitoring strategy, it is important to determine not only (a) which patients are at highest risk, but also (b) how long monitoring should take place; yet we have limited data available for learning because the outcome rates are low.

Our goal, therefore, is to develop a principled, data-driven approach to answer both of these questions. Specifically, we wish to develop a method that providers can use to identify time intervals that are optimal when understanding risk, for example to design an intervention or monitoring strategy, as well as when reporting risk to patients. At the same time, we wish to retain the substantial advantages and flexibility of other recently developed approaches, including their lack of strong parametric assumptions about the form of the event density. To solve this problem, we develop Adaptive Discretization for Event PredicTion (ADEPT).

We begin by recasting learning from discrete survival times as learning from continuous survival times under the assumption that the density is piecewise constant; and then formulate a smooth relaxation of this piecewise constant density that allows cut points (i.e., interval boundaries) to be learned by gradient-based optimization methods. We then present our learning procedure and results of experiments with two simulated and three real datasets – including a newly harmonized stroke risk prediction dataset that pools data across three large cohorts – that illustrate the effectiveness and potential clinical relevance of ADEPT.

Our performance evaluation focuses on comparing our method to its state of the art alternative, namely, discrete-time, neural network-based risk prediction over fixed-length intervals.

In summary, our contributions are as follows:

• Present ADEPT, a novel model and associated learning procedure to learn an optimal event time partition from data rather than fixing it a priori.

• Present simulation results illustrating effective learning of cut points that are consistent with the true, underlying generative model.

• Demonstrate improved prediction performance across three real datasets, including two clinical datasets.

• Identify clinically meaningful risk cut points illustrating the potential of the approach to provide improved prognostic information.

2. METHODS

2.1. Setup and Notation

Consider a time-to-event outcome where each observation is represented by the triplet $\{\mathit{X},Y,S\}$, where $\mathit{X}\in 𝒳\in {\mathbb{R}}^p$ is a $p$-dimensional feature vector, $Y\in \left(0,T_{\mathrm{m}\mathrm{a}\mathrm{x}}\right]$ is an observed event time over a finite time horizon, and $S\in \{0,1\}$ indicates whether $Y$ is a right-censoring time $(S=0)$ or an event time $(S=1)$. The observed time $Y$ is the minimum of the event time $T$ and the right-censoring time $U$, i.e., $Y=\mathrm{m}\mathrm{i}\mathrm{n}(T,U)$, and $S=\mathbb{I}(T<U)$, where the indicator function $\mathbb{I}(\cdot )$ is 1 when the argument is true and 0 otherwise.

We consider possible sequences of $M$ cut points $C={\left\{c_j\right\}}_{j=1}^M$, where $0=c_0<c_1<\cdots <c_M<c_{M+1}=T_{\mathrm{m}\mathrm{a}\mathrm{x}}$, that partition the event time space, $\left(0,T_{\mathrm{m}\mathrm{a}\mathrm{x}}\right]$, into the intervals $I_1,\dots ,I_{M+1}$, where $I_j=\left(c_{j-1},c_j\right]$. Figure 1 provides an example of the event time space partitioned into four intervals: $I_1=\left(0,c_1\right],I_2=\left(c_1,c_2\right],I_3=\left(c_2,c_3\right]$, and $I_4=\left(c_3,T_{\mathrm{m}\mathrm{a}\mathrm{x}}\right]$. Given such a partition, we introduce an auxiliary random variable $Z\in \{1,\dots ,M+1\}$ that indicates which interval contains $T$, i.e., $Z=j⟺t\in I_j$.

Figure 1:

The event time space partitioned by three cut points into four intervals.

2.2. Piecewise Constant Density

We begin by considering learning with fixed cut points, which is currently the predominant approach. For example, Lee et al. [2018] and other recently-developed methods [Ren et al., 2019, Tjandra et al., 2021, Engelhard and Henao, 2022] use fixed cut points to discretize time in order to avoid placing restrictive, parametric assumptions on the form of the event time density. Instead, the density is restricted to be piecewise constant according to the intervals defined by the cut points. The cut points themselves might be evenly spaced in time, or alternatively they might be evenly spaced across the observed or estimated event time distribution, e.g., via empirical quantiles. The goal of learning is then to estimate $P(Z\mid \mathit{X})$, the conditional probability that $T$ will fall in each of the pre-defined intervals, rather than $p(T\mid \mathit{X})$, the conditional density of $T$. Typically $T$ is discretized to $Z$ a priori.

However, it is not possible to learn the cut points $C$ with this approach, because $Z$ depends on $C$ in addition to $T$. To see this, consider the value of $Z$ associated with an observed time $t\in \left(0,T_{\mathrm{m}\mathrm{a}\mathrm{x}}\right]$ under the binary partition defined by the single cut point $c_1$. If we choose $c_1\ge t$, then $t\in \left(0,c_1\right]$, therefore $Z=1$; but for $c_1<t$, we have $t\in \left(c_1,T_{\mathrm{m}\mathrm{a}\mathrm{x}}\right]$, therefore $Z=2$.

To circumvent this limitation, we note that estimating $P(Z\mid \mathit{X})$ is equivalent to estimating $p(T\mid \mathit{X})$ with the following piecewise constant model, which supposes $p(T\mid \mathit{X})$ has uniform density over each interval $I_j$:

$$\widehat{p}\left(t\mid x\right)=\sum _{j=1}^{M+1}{p}_{\varphi }\left({𝓏}_j\mid x\right)\frac{{\mathbb{I}}_{I_j}\left(t\right)}{\left|I_j\right|},$$ (1)

where ${\mathbb{I}}_{I_j}(\cdot )$ is the indicator function associated with the interval $I_j$, and $\varphi$ parameterizes our model of $P(Z\mid \mathit{X})$. Importantly, we must normalize by $\left|I_j\right|$, the length of $I_j$, to ensure ${\int }_{\left(0,T_{\mathrm{m}\mathrm{a}\mathrm{x}}\right]}\hat{p}(t\mid \mathit{x})=1$ and ${\int }_{I_j}\hat{p}(t\mid \mathit{x})=p_{\varphi }\left(z_j\mid \mathit{x}\right)$. As a potential drawback of this approach, we note that whereas standard discrete-time approaches are well suited to handle outlying event times, in principle this normalization term could cause the loss to become numerically unstable in the case of extreme event time outliers.

2.3. Smooth Relaxation of Piecewise Density

The parameters $\varphi$ of our model for $Z$ can be learned directly from equation (1). However, our goal is to learn not only $\varphi$ but also $C$, the specific partition that allows our model to best approximate $p(T\mid \mathit{X})$ across a given dataset. Unfortunately, (1) cannot be optimized with respect to $C$ via gradient-based methods. This is because the indicator function ${\mathbb{I}}_{I_j}(\cdot )$ implicitly depends on $C$, and is discontinuous whenever a cut point is equal to an observed event time.

To illustrate, consider learning a single cut point $c_1$ while holding the parameters $\varphi$ fixed. For small $\epsilon$ such that $0<\epsilon <t$, where $t$ is an observed event time associated with covariates $\mathit{x}$, suppose the cut point $c_1=t+\epsilon$ is just after the observed event time. In this case, we have $t\in I_1$, therefore ${\mathbb{I}}_{I_1}\left(t\right)=1$ and ${\mathbb{I}}_{I_2}\left(t\right)=0$, and consequently $\hat{p}(t\mid \mathit{x})=p_{\varphi }\left(z_1\mid \mathit{x}\right)/\left|I_1\right|$. On the other hand, suppose the cut point $c_1=t-\epsilon$ is just before the observed event time. In this case, we have $t\in I_2$, therefore ${\mathbb{I}}_{I_1}\left(t\right)=0$ and ${\mathbb{I}}_{I_2}\left(t\right)=1$, and consequently, $\hat{p}(t\mid \mathit{x})=p_{\varphi }\left(z_2\mid \mathit{x}\right)/\left|I_2\right|$. Thus, for any non-trivial model $p_{\varphi }$ for which $p_{\varphi }\left(z_1\mid \mathit{x}\right)\ne p_{\varphi }\left(z_2\mid \mathit{x}\right)$, equation (1) is discontinuous at $c_1=t$. This argument readily generalizes to all cut points.

To smooth this discontinuity and allow gradient-based optimization, we replace the indicator function ${\mathbb{I}}_{I_j}\left(t\right)$ in (1) with the smooth approximation $\sigma \left(\left(t-c_{j-1}\right)/\tau \right)*\sigma \left(\left(c_j-t\right)/\tau \right)$, where $\sigma \left(z\right)={\left(1+e^{-z}\right)}^{-1}$ is the sigmoid function. The temperature $\tau$ is a hyperparameter of the model that should be tuned based on the scale of the observed event times.

This results in the following relaxed model:

$$\widehat{p}\left(t\mid x\right)=\sum _{j=1}^{M+1}{p}_{\varphi }\left({𝓏}_j\mid x\right)\frac{\sigma \left(\frac{t-c_{j-1}}{\tau }\right)\sigma \left(\frac{c_j-t}{\tau }\right)}{\left|I_j\right|},$$ (2)

which is approximately piecewise constant for $\tau \ll T_{\mathrm{m}\mathrm{a}\mathrm{x}}$, yet differentiable everywhere with respect to $C$ and thus suitable for gradient-based optimization.

2.4. Learning Procedure

Under the common assumption of non-informative right-censoring, we may ignore the censoring density and optimize $\hat{p}(y,s\mid \mathit{x};\theta )$, where $\theta =\{\varphi ,C\}$, over the observed data $𝒟={\left\{{\mathit{x}}_i,y_i,s_i\right\}}_{i=1}^N$ as follows:

$$\theta =\underset{\theta }{\text{arg}\text{max}}\sum _i^N\left[s_i\text{log}\widehat{p}\left(t_i\mid \mathit{x};\theta \right)\right.\left.+\left(1-s_i\right)\text{log}\widehat{P}\left(t_i>y_i\mid {\mathit{x}}_i;\theta \right)\right],$$ (3)

where $\hat{P}\left(t_i>y_i\mid {\mathit{x}}_i;\theta \right)=1-{\int }_0^T\hat{p}\left(\tau \mid {\mathit{x}}_i;\theta \right)$ is the survival function associated with $\hat{p}\left(t_i\mid {\mathit{x}}_i;\theta \right)$ for observation $i$.

However, optimizing equation (2) alone can result in degenerate solutions in which cut points become arbitrarily close together or even coincide. In the extreme case, it is possible to have $I_j=\left(0,T_{\text{max}}\right]$ for a particular $j\in \{1,M+1\}$, resulting in the trivial model in which $p_{\varphi }(z\mid \mathit{x})$ places all mass on $z_j$.

It is therefore critical to balance optimizing equation (2) versus ensuring that $p_{\varphi }(z\mid \mathit{x})$ is non-trivial. We accomplish this by incorporating a regularization term, $H\left(p_{\varphi }(z\mid \mathit{x})\right)$, with associated hyperparameter ${\lambda }_1$ in our optimization procedure. We use a scaled Beta (1.5, 1.5) distribution on each cut point. For example, suppose there are three cut points $c_1<c_2^{\star }<c_3$ where $c_2^{\star }$ is the newly proposed value for the middle cut point $c_2$. We scale the value of the cut point to find its location relative to the cut points near it: $c_{2,\text{scaled}}^{\star }=\left(c_2^{\star }-c_1\right)/\left(c_3-c_1\right)$. The final regularization value is the PDF value of $c_{2,\text{scaled}}^{\star }$ evaluated over a Beta (1.5, 1.5) distribution. This regularization term encourages cut points to be near the center of their two surrounding cut points.

We may then optimize $\theta$ over $𝒟$ by choosing $\theta =\text{arg}{\mathrm{m}\mathrm{i}\mathrm{n}}_{\theta }{\sum }_{𝒟}ℒ\left(\theta \right)$, where $ℒ\left(\theta \right)$ is defined as follows:

$$ℒ\left(\theta \right)=-\text{log}\widehat{p}\left(y,s\mid x;\theta \right)-{\lambda }_{1}H\left(p_{\varphi }\right).$$ (4)

Here the first term is the negative log likelihood in (3) and the second is our beta distribution regularizer. Our learning procedure then becomes:

$$\theta =\underset{\theta }{\text{arg}\mathrm{m}\mathrm{i}\mathrm{n}}\sum _{𝒟}ℒ\left(\theta \right)+{\lambda }_{2}R\left(\theta \right),$$ (5)

where we have included an additional regularization term $R\left(\cdot \right)$ (e.g., $L_2$-regularization) along with an associated hyperparameter ${\lambda }_2$ to control for overfitting.

3. IMPLEMENTATION DETAILS

3.1. Baseline Model: Discrete-Time Neural Network

We compare ADEPT to a discrete-time neural network baseline that is identical to the proposed model, except the cut points (and corresponding intervals) are initialized based on the observed outcomes and remain fixed when learning the classifier. This approach, hereafter called the DTNN, was popularized by DeepHit [Lee et al., 2018] and is currently the predominant approach. However, to isolate the effect of learning the partition we do not include the ranking loss used in DeepHit. Our method and the DTNN have similar computational complexity, which is dominated by the computation of gradients with respect to model parameters $\theta$ rather than the cut points $C$. In this work, we instantiate $p_{\theta }$ as neural network, but our approach is flexible to the model choice; thus it can be changed based on application if, for example, interpretability is more important that predictive performance.

We initialize the DTNN model’s cut points to be evenly spaced on the percentiles of the empirical Kaplan-Meier curve of the observed outcomes; for example, if there are three cut points then they would be placed at the time points associated with the 25^th, 50^th, and 75^th percentiles on the estimated Kaplan-Meier curve. With these cut points fixed, we then build a model predicting the probability that the patient will experience the outcome in each interval. Note that this differs from ADEPT in which we also consider the cut points themselves as parameters. The DTNN classification model learns the probability of each observation being in each of the pre-defined intervals. In the notation of Section 2.2, the DTNN approach learns only the model parameters, ${\varphi }_{\text{baseline}}$, whereas ADEPT learns both model parameters $\varphi$ and the cut points $C$. Importantly, we search the same grid of hyperparameters for the DTNN model as for ADEPT.

Due to its popularity, we also include a comparison to DeepSurv [Katzman et al., 2018]. However, understanding differences between DeepSurv and other approaches is challenging due to their stark differences, notably the assumption of proportional hazards, which can either improve or worsen performance depending on the degree to which this assumption holds. Moreover, because DeepSurv does not incorporate discretization, it will have fewer comparative performance metrics.

3.2. Performance Quantification

With simulated data we were able to judge the correctness of the estimated cut points by their proximity to the true cut points used in the data generation mechanism. For the real data we do not know the true cut point values and thus need other metrics to judge performance. While we focus on several metrics quantifying predictive performance, clinical collaboration is necessary to determine which metrics are most relevant in a particular clinical context, including when implementing treatment decisions.

Time-Dependent Concordance Index (CI)

Since we consider a time-to-event outcome with censored observations rather than a regression or classification outcome, standard metrics such as root mean square error and area under the receiver operating characteristic are insufficient to capture prediction performance. Initially developed by Harrell Jr et al. [1984], the concordance index (CI) measures how well predicted event times match the order of the true event times. However, both ADEPT and the DTNN predict discrete interval membership rather than continuous event times, and the ordering of predicted risk can change over time. To properly account for these characteristics, we use a discrete-time implementation of the time-dependent concordance index developed by Antolini et al. [2005]. This metric compares model-predicted risk at observed failure times to the model-predicted risks at that time for other individuals known to have later failure times. Pairs of individuals are only considered if (a) both failure times are known (neither are censored), or (b) one failure time is known to have occurred before the censoring time of the other.

AUC at last cut point

The Area Under the Receiver Operating Characteristic Curve (AUC) is a common metric to evaluate predictive performance for a binary outcome. To adapt this to our method, we focus on the AUC at the last cut point. That is, we are interested in determining if the methods are able to predict whether an event happens before or after the final cut point. This is especially relevant for data sets with high amounts of censoring at the end of the study. The cases are all observations that experienced an event prior to the final cut point and the controls are all observations with an observed time (either an event or censored) after the final cut point. Notice that observations that are censored prior to the last cut point are omitted from this metric.

Integrated Brier Score (IBS)

The Brier Score evaluated at time a chosen time $t$ is the mean squared difference between the model-predicted cumulative event probability at time $t$ and the true, binary, observation of whether the even occurred by $t$. The Integrated Brier Score improves upon this by integrating over all times $t\in \left\{t_{\mathrm{m}\mathrm{i}\mathrm{n}},t_{\mathrm{m}\mathrm{a}\mathrm{x}}\right\}$ [Graf et al., 1999]. An IBS of 0 indicates that the model was able to perfectly predict outcomes.

Calibration slope and intercept

We also consider the calibration slope and intercept as described by Crowson et al. [2016], which quantify the degree to which model-predicted probabilities accurately estimate true event probabilities, as determined based on observed event rates. A well calibrated model will have a calibration slope near 1 and a calibration intercept near 0.

It is important to note that DeepSurv does not discretize the event time space; because of this, only the AUC and IBS metrics are included. We cannot compare the discrete-time concordance index to the standard concordance index because there are inevitably more ties with the discrete-time approach.

3.3. Hyperparameter Tuning

ADEPT is flexible, allowing for any number of cut points. In our simulation examples we will know exactly how many cut points were used to generate the data; however, this is not the case for the real data experiments. Thus, we use 3, 5, and 10 cut points. We use a two layer neural network as our predictive model.

The first layer has input dimension $p$ based on the feature dimension of the data and output dimension $h$, for which we explore values of 32, 128, and 512. This is then connected by a Rectified Linear Unit activation function to another layer with input dimension $h$ and output dimension. These networks are optimized using Adam [Kingma and Ba, 2014] with a learning rate of 0.01 and weight decay values between 0.0001 and 0.1. We vary the strength of the regularization on the cut points ${\lambda }_1$ from values in the range of 0.1 to 20 and use a mini-batch size of 64 for the training data.

During training, we initially set the sigmoid temperature used in our smooth approximation (see Equation (2)) to a value $\tau =0.1$, then lower it when the loss stops changing significantly between epochs. Lowering the temperature reduces the degree of smoothing and sharpens the boundaries between intervals defined by each cut point. Figure 2 shows an example training plot where the temperature drops after multiple epochs with no improvement in the validation loss. It is clear that this drop then leads to an improvement in both training and validation loss.

Figure 2:

A training plot of the training and validation loss at each epoch for the two interval simulation example. The vertical lines represent drops in sigmoid temperature $\tau$ and the accompanying new value of $\tau$.

We perform a grid search over the hyper parameters, testing every combination of output dimension, weight decay, and regularization strength. The evaluation process to compare hyperparameters is described in Section 3.2. We train each network for 250 epochs.

To evaluate performance we perform five-fold cross validation. We randomly partition the data into training (75%), validation (15%), and test (10%) sets. For each set of hyperparameters we perform this partition five times, using the training sets for learning the model parameters. We then calculate average performance metrics metrics on the out of sample validation sets.

Only the model with the best average validation set performance is then applied to the corresponding, yet unseen, test sets. We report the average and standard deviations of the performance metrics calculated across the folds on the test sets. Through this general cross-validation strategy, we are able to find the hyperparameter setting that performs the best on out of sample data from the hyperparameters tested. We report the mean and standard deviation of each metric across the folds.

We perform the same parameter search and evaluation to find the best DTNN model as described in Section 3.1. We compare the metrics of the best ADEPT model to that of the best baseline models. We report the performance metrics all both methods calculated on the unseen test set.

4. SIMULATION EXAMPLES

4.1. Learning Two Intervals

We start with the simple case of data generated from two clusters with uniform censoring. Cluster membership is generated using the make_moons function sklearn Python package to get a noisy, nonlinear relationship between $p=2$ features [Pedregosa et al., 2011]. Figure 3a shows the feature-cluster relationship; each cluster has 5, 000 observations for a total of $n=7,500$ observations. These clusters are used to generate the event times.

Figure 3:

The event times and observed times of the two interval data. The true cut point is at time 67.

Event times in Cluster 1 are generated uniformly on the interval (0, 67] and event times in Cluster 2 are uniformly on the interval (67, 100). Censoring times are then generated uniformly throughout the entirety of (0, 100). Note that while the censoring and event times are both uniformly distributed, the observed times are the minimum of the two and therefore not uniformly distributed. These observed times are shown in Figure 3b. Because these intervals are determined by the relationship between the covariates, this simulates data generated with a true cut point at time 67.

Figure 3c shows out of sample test set along with the DTNN cut point in red at $t=49.3$ and the ADEPT learned cut point in black at $t=65.6$. Knowing that the true cut point is at time 67 demonstrates the efficacy of our method. Even with a starting point far from the true cut points, we are able to recover the true cut point. Table 1 reports the performance metrics, showing a large gain in CI.

Table 1:

Performance metrics for synthetic data. We report average metrics across 5 -fold cross validation with standard errors in parentheses. In bold are the highest CI values for each setting.

Intervals	ADEPT CI	DTNN CI
Two (n = 7, 500)	0.947 (0.001)	0.797 (0.002)
Two (n = 300)	0.813 (0.039)	0.756 (0.016)
Four (n = 7, 500)	0.980 (0.012)	0.937 (0.007)
Four (n = 300)	0.964 (0.013)	0.931 (0.013)

In this simple example, many combinations of hyperparameters were able to recover the true cut point; reported are the results from using a small neural network with $h=32$ with Adam weight decay of 0 and a regularization strength of $\lambda =1$.

To demonstrate that our method works in limited data settings, we repeat this experiment randomly selecting only $n=300$ training data points. Despite this data limitation, the blue lines in Figure 3c shows that ADEPT was still able to recover the true cut point. Table 1 shows that ADEPT still outperforms the DTNN model in prediction.

4.2. Learning Four Intervals

With confidence in ADEPT’s ability to learn a single cut point when it is present in the data generation, we expand to learning three true cut points. Again we use the make moons function to generate noisy, nonlinear relationships between $p=2$ features, however now for four separate clusters as shown in Figure 4a; each cluster has 2, 500 observations for a total of $n=7,500$ observations. Figure 4b shows the how these clusters are used to generate event times. Event times are generated using a Beta (1.5, 1.5) distribution which are then scaled to be in the appropriate interval based on the observation’s cluster. The first cluster has event times on the interval (0, 10], the second on the interval (10, 30], the third on the interval (30, 70], and the fourth on the interval (70, 100). This corresponds to the true but points being at $t=10,30,70$. We again apply uniform censoring times to all observations. Note that with uniform censoring, there are particularly few uncensored observations for events in the last interval. This makes learning the final cut point more difficult.

Figure 4:

Event times and observed times of the four interval data. The true cut points are at 10, 30, and 70.

The black lines in Figure 4c shows that ADEPT was able to successfully recover all three cut points despite the challenges due to censoring. Table 1 shows that the ADEPT’s learned intervals provide an increase in CI over the DTNN. Since we know the data generating mechanism, it is intuitive for this simulation example that including more than 3 cut points leads to worse performance as introducing more would over-parametrize the model. The results in the next section suggest that it is beneficial to consider models with fewer cut points even when the generating mechanism is unknown.

We repeat this simulation with only $n=300$ data points. With the uniform censoring, this results in even fewer observed events in the final cluster. As in the two interval case, ADEPT is still able to recover the true cut points, shown in Figure 4c, and outperform the DTNN baseline model in predictive performance, shown in Table 1.

5. DATA ANALYSIS

5.1. Real-World Data Sources

We apply our method to three real-world data sources of varying sizes.

German Breast Cancer Study Group (GBSG)

The GBSG data set is a publicly available data set introduced by Schumacher et al. [1994]. It is a multicenter clinical trial which includes $n=686$ patients with $p=8$ features. The endpoint of recurrence free survival occurred for 299 (43.6%) patients.

Assay of Serum Free Light Chain (FL Chain)

The FL Chain data set is a publicly available data set introduced by Dispenzieri et al. [2012] studies the relationship between nonclonal serum immunoglobulin free light chains and mortality. We examine the data for the $n=6,524$ patients that had no missing data with $p=8$ features. The endpoint of death occurred for 1, 962 (30.1%) of these patients.

Pooled Stroke Risk Cohorts

This is a combined dataset consisting of the Framingham Offspring Study [Feinleib et al., 1975] $\left(n_1=8,348\right)$, The Atherosclerosis Risk in Communities Study [Investigators, 1989] $\left(n_2=23,158\right)$, and the Multi-Ethnic Study of Atherosclerosis $\left(n_3=6,390\right)$ [Bild et al., 2002]. Data harmonization procedures and characteristics of the dataset have previously been described by Hong et al. [2023]. We consider a total of $n=35,450$ data points of which 1, 221 (3.44%) experience a stroke. There are $p=69$ features that include cardiovascular medical history, demographic indicators, and diet information.

5.2. Results

Figure 5 shows the best ADEPT learned cut points for each data set compared to the DTNN. Note Figure 5c is a histogram of the proportion of observations rather than raw counts because of the high amount of censoring. Table 2 shows the performance metrics for all real-world data sets. The reported metrics are calculated on the held-out test sets not used for training or model validation.

Figure 5:

The DTNN (red, dashed) and ADEPT learned (black, solid) cut points.

Table 2:

Test-set performance metrics for real-world data. Reported are average metrics across 5-fold cross validation with corresponding standard errors in parentheses. In bold are the highest CI, highest AUC, and lowest IBS models for each data set.

	3 Cut Points	5 Cut Points	10 Cut Points
GBSG
ADEPT CI	0.744 (0.015)	0.68 (0.018)	0.671 (0.024)
DTNN CI	0.681 (0.027)	0.651 (0.059)	0.619 (0.065)
ADEPT AUC	0.804 (0.021)	0.801 (0.02)	0.822 (0.024)
DTNN AUC	0.800 (0.03)	0.750 (0.034)	0.807 (0.016)
DeepSurv AUC	0.795 (0.054)	—	—
ADEPT IBS	0.180 (0.005)	0.189 (0.012)	0.173 (0.006)
DTNN IBS	0.187 (0.006)	0.198 (0.008)	0.194 (0.005)
DeepSurv IBS	0.165 (0.007)	—	—
ADEPT Calibration Slope	0.813 (0.089)	0.995 (0.112)	0.793 (0.165)
DTNN Calibration Slope	1.00 (0.091)	1.855 (0.154)	1.397 (0.254)
ADEPT Calibration Intercept	0.178 (0.048)	0.130 (0.042)	0.142 (0.06)
DTNN Calibration Intercept	0.129 (0.057)	−0.285 (0.084)	−0.215 (0.129)
FL Chain
ADEPT CI	0.798 (0.003)	0.793 (0.003)	0.787 (0.004)
DTNN CI	0.763 (0.007)	0.772 (0.004)	0.774 (0.012)
ADEPT AUC	0.806 (0.004)	0.81 (0.004)	0.834 (0.002)
DTNN AUC	0.788 (0.008)	0.809 (0.004)	0.828 (0.004)
DeepSurv AUC	0.831 (0.002)	—	—
ADEPT IBS	0.194 (0.005)	0.163 (0.006)	0.146 (0.005)
DTNN IBS	0.191 (0.008)	0.153 (0.004)	0.13 (0.002)
DeepSurv IBS	0.099 (0.001)
ADEPT Calibration Slope	1.199 (0.103)	1.182 (0.075)	1.057 (0.098)
DTNN Calibration Slope	1.005 (0.04)	0.999 (0.039)	0.898 (0.076)
ADEPT Calibration Intercept	0.056 (0.007)	0.050 (0.01)	0.085 (0.018)
DTNN Calibration Intercept	0.102 (0.008)	0.100 (0.009)	0.116 (0.026)
Stroke
ADEPT CI	0.789 (0.014)	0.747 (0.02)	0.765 (0.006)
DTNN CI	0.778 (0.01)	0.739 (0.017)	0.758 (0.022)
ADEPT AUC	0.766 (0.011)	0.701 (0.03)	0.723 (0.013)
DTNN AUC	0.743 (0.01)	0.681 (0.031)	0.713 (0.021)
DeepSurv AUC	0.705 (0.015)
ADEPT IBS	0.027 (0.002)	0.029 (0.002)	0.03 (0.002)
DTNN IBS	0.032 (0.003)	0.033 (0.002)	0.032 (0.002)
DeepSurv IBS	0.019 (0.001)
ADEPT Calibration Slope	0.783 (0.013)	1.117 (0.195)	1.270 (0.321)
DTNN Calibration Slope	1.098 (0.321)	1.370 (0.248)	1.295 (0.376)
ADEPT Calibration Intercept	0.019 (0.002)	0.022 (0.003)	0.025 (0.002)
DTNN Calibration Intercept	0.022 (0.003)	0.019 (0.002)	0.025 (0.002)

Two interesting trends strongly support the benefits of ADEPT. First, for all data sets, the CI was the highest for the models that used only 3 cut points and tended to decrease as more cut points were added. Additionally, for all numbers of cut points, the predictive performance in both CI and AUC for the learned cut point model was higher than the DTNN model.

Notice that the greatest improvement in CI was observed for the GBSG data set, which has the fewest observations among all data sets. This underscores the importance of ADEPT. For small data sets with a limited number of outcomes, it is necessary to limit the number of cut points, but performance can be improved by optimizing their locations. Notice that for the FL Chain data set, the DTNN model achieved its highest CI with 10 cut points, but this was still lower than the performance of ADEPT using only 3 learned cut points. Interestingly, the Stroke data set, which had the most data points and the highest outcome imbalance, also had a higher CI and AUC with 10 cut points than with 5. Similar to the other data sets, it achieved its highest CI and AUC using 3 cut points.

A model that is well calibrated has a calibration slope near 1 and a calibration intercept near 0. While the DTNN model had slightly better calibration slopes for the GBSG and FL Chain data sets for 3 cut points, ADEPT was better calibrated in nearly every setting with more cut points demonstrating model robustness.

While ADEPT was able to outperform DTNN in IBS for any given number of cut points, the continuous prediction of DeepSurv had the lowest overall IBS for each data set.

6. CONCLUSION

Herein we develop ADEPT, a flexible method to learn an optimal partitioning of the event time space that does not place strong assumptions on the form of the event density. Our approach is designed for clinical applications in which it is advantageous to learn, from data, a time discretization that facilitates more accurate prediction. The simulated examples demonstrated the ability of our method to recover cut points when they are truly present in the data generation mechanism. Moreover, results on real data show that the approach improves prediction performance over otherwise equivalent, state of the art models that use a fixed discretization scheme.

Our approach can be extended in several ways. In future work, we will consider a similar approach to learn separating hyperplanes in higher dimensional output spaces. The method can also be extended to sequential or time series data by using an appropriate encoder (e.g., a recurrent neural network). Another interesting extension motivated by the real-world data analysis would be to learn the number of cut points from the data instead of fixing it a priori.