Censored Quantile Regression Neural Networks for Distribution-Free Survival Analysis
Abstract
This paper considers doing quantile regression on censored data using neural networks (NNs).
This adds to the survival analysis toolkit by allowing direct prediction of the target variable, along with a distribution-free characterisation of uncertainty, using a ﬂexible function approximator.
We begin by showing how an algorithm popular in linear models can be applied to NNs.
However, the resulting procedure is inefﬁcient, requiring sequential optimisation of an individual
NN at each desired quantile.
Our major contribution is a novel algorithm that simultaneously optimises a grid of quantiles output by a single NN.
To offer theoretical insight into our algorithm, we show ﬁrstly that it can be interpreted as a form of expectation-maximisation, and secondly that it exhibits a desirable ‘self-correcting’ property.
Experimentally, the algorithm produces quantiles that are better calibrated than existing methods on 10 out of 12 real datasets.
Introduction
Domains such as biomedical sciences and reliability engineering often produce datasets with a particular challenge – for many datapoints the target variable is not directly observed and only a lower or upper bound is recorded.
For example, when modelling the time to failure of a machine, in many cases it might have been preemptively maintained before a failure event was observed, meaning only a lower bound on the time to failure is recorded.
This is known as censored data, which is studied in the ﬁeld of survival analysis.
Recently, there has been interest in combining ideas from survival analysis with neural networks (NNs), in the hope of leveraging the capabilities of deep learning for this important class of problem.
Quantile regression has proven valuable in survival analysis, since it directly predicts the variable of interest, naturally capturing the uncertainty of the conditional distribution with no assumptions made about the distribution of residual errors [Peng, 2021].
Despite a wealth of research in the linear setting, explorations into its combination with NNs are at an early stage [Jia and Jeong, 2022].
This paper advances this line of research, showing how an estimator from Portnoy [2003], used for doing quantile regression in linear models on left and right censored data, can be combined with NNs.
Section 2 introduces Portnoy’s estimator, and outlines a sequential grid algorithm that is used in its optimisation for linear models [Neocleous et al., 2006].
As our ﬁrst contribution, Section 3.1 shows how this algorithm can be directly adapted to work with NNs.
Unfortunately, the resulting method is inefﬁcient, requiring sequential optimisation of a new NN at each predicted quantile.
Section 3. outlines our main contribution; an improved algorithm that is well suited to NNs, named ‘censored quantile regression neural network’ (CQRNN).
Our CQRNN algorithm is a signiﬁcant departure from the sequential approach, and we offer theoretical insights into it in Section 4, interpreting it as a form of expectation-maximisation (EM), and also analysing a ‘self-correcting’ property.
In Section 6 the algorithm’s effectiveness is empirically demonstrated by benchmarking against alternative methods on simulated and real data.
Careful ablations illuminate effects of the algorithm’s hyperparameters.
Background
We ﬁrst introduce censored data.
Consider a dataset D = {{ x 1 , y 1 ,  1 } , ..., { x N , y N ,  N }} , where x i 2 R D is the input, y i 2 R is the (possibly censored) variable to regress on, and  i represents the observed/censored indicator.
We assume two data generating distributions, one for the target variable, t i ⇠ p t ( t | x i ) , and another for the censoring variable, c i ⇠ p c ( c | x i ) .
In the case of right censoring, we only observe the smaller of these two, which forms our dataset labels, y i = min( t i , c i ) , and  i = [-(cid:110) 
-]n 0 if censored, c
i <t i 1 else —————— .
Whilst this paper concentrates on the right censored case, for all estimators and algorithms discussed, left censoring can be handled by simply inverting the labels, i.e. using {+ y i 
+}
i [Koenker, 2022].
In this paper we make common assumptions about p t ( t | x i ) and p c ( c | x i ) ; the target is independent of censoring given the covariates, t ? c | x (sometimes termed ‘random’ censoring), but both t and c depend on x in different and possibly non-linear ways.
These are more general than some alternative assumptions, such as using ﬁxed-value censoring, c i =
constant 8 i (e.g. [Powell, 1986]) or requiring the censoring distribution to be independent of x (e.g. [Jia and Jeong, 2022]).
2.1 Quantile Regression without Censoring
Ignoring censoring for a moment, the conditional quantile function is given by,
[Equation]
where, ⌧ 2 (0 , 1) .
To learn such a function for a single target quantile, ⌧ , quantile regression minimises a ‘checkmark’ loss, ⇢ ⌧ ( · , · ) : R ⇥ R !
R , [Koenker and Bassett, 1978],
[Equation]
for model parameters, ✓ , and a prediction, ˆ y i, ⌧ = ⌧ ( x i , ✓ ) , made by some model ⌧ : R D !
R , with I [ · ] as the indicator function.
A linear model refers to when, ˆ y i, ⌧ = ✓ | ⌧ x i , with, ✓ ⌧ 2 R D .
2.2 Portnoy’s Censored Quantile Regression Estimator
Directly optimising the loss in Eq.
2 can produce undesirable models when a dataset contains censored observations.
Naively ignoring censoring indicators or excluding all censored data will, in the general case, induce unfavourable bias [Zhong et al., 2021].
The algorithms proposed in this paper use a re-weighting scheme introduced in Portnoy [2003] that does account for censoring.
Deﬁne two disjoint sets of indices, one for censored and one for observed datapoints, S censored and S observed , letting, N c :
=
|S censored | and N o :=
|S observed | .
Portnoy’s estimator minimises,
L Port.
( ✓ , D , ⌧ , w , y ⇤
[Equation]
=
[Equation]
While the loss for observed datapoints is unchanged, censored datapoints have been split into two ‘pseudo’ datapoints – one at the censored value, and one at some large value, y ⇤  max i y i
(Section 6.4 discusses y ⇤ further).
Weight is apportioned between each pair of pseudo datapoints by,
[Equation]
where q j is the quantile at which the datapoint was censored, with respect to the target value distribution, i.e. p t ( t j > c j | x j ) .
We use w 2 R N c and q 2 R N c to denote the vector of all weights and quantiles respectively, indexed as w j and q j .
Given these weights, Portnoy’s estimator has been shown to be analogous to the Kaplan-Meier (KM) estimator
[Portnoy, 2003].
Require: Dataset D , M parametric models ⌧ ( · ) each with randomly initialised parameters ✓ ⌧ , ordered quantiles to be estimated grid ⌧ , learning rate ↵ , pseudo y value y ⇤ .
S censored { i 2 { 0 , 1 , . . .
, N } :  i = 0 } , S observed { i 2 { 0 , 1 , . . .
, N } :  i = 1 } , K cross ; for i = 0 to M  1 do ⌧ grid ⌧ [ i ] if i = 0 then .
Initialise quantile estimates to zero ˆq zeros(N c ) else .
Find indices which have been crossed and update their estimated quantiles
[Equation]
Sets ˆw = 0 for uncrossed datapoints
[Equation]
Fully optimise Eq.
2.3 Sequential Grid Algorithm
A challenge with Portnoy’s estimator is that it creates a circularity.
If the true quantiles q and hence w for all censored datapoints were known, Eq. 4 could be optimised straightforwardly, but the very reason to perform this optimisation is to estimate such quantiles!
Prior work has tackled this issue in various ways (see Section 5).
Here we discuss the most widely-used algorithm.
It was originally presented in Portnoy [2003] and slightly modiﬁed in Neocleous et al.
[2006], we refer to it as the ‘sequential grid algorithm’.
Described in Algorithm 1, it sequentially steps through a grid of M desired quantiles, arranged in strictly increasing order, and typically evenly spaced, e.g., grid ⌧ = { 0 . 1 , 0 . 2 , . . .
, 0 . 8 , 0 . 9 } for
M = 9 .
A new model is ﬁtted at each quantile.
The algorithm terminates either when all quantiles in the grid have been iterated through, or if only censored datapoints lie above the current quantile.
(Portnoy [2003] suggest handling the ﬁrst quantile specially, but we have simpliﬁed this – see later.)
Note we introduce ˆq & ˆw to explicitly designate model estimates of q & w .
[Figure]
N o r m un
i f
o r m
[Figure]
N o r m li n ea r
E xpon e n ti a l
L og
N o r m
This section considers applying the sequential grid algorithm to NNs, explaining why this is inefﬁcient.
It then describes the CQRNN algorithm, our main contribution, which avoids these pitfalls.
3.1 The Sequential Grid Algorithm and Neural Networks
To the best of our knowledge the sequential grid algorithm has been proposed and implemented only on linear models.
Our ﬁrst contribution is showing that, with three adjustments, it can be used in NNs.
Adj. 1)
The M models producing estimates, ˆ y j, ⌧ , are chosen be NNs, of any architecture with a single output.
Adj. 2)
While the linear version of the algorithm minimises Eq.
4 analytically with a parametric programming approach, for NNs this is instead minimised with stochastic gradient descent.
Adj. 3)
Portnoy [2003] proposed using a lengthy process for the ﬁrst quantile.
They set ˆw = 1 and optimise Eq. 4.
After optimisation, if any censored datapoints lie below the ﬁrst quantile, these are removed from the dataset, and the procedure repeats until there are none.
This repeated optimisation could go on for many iterations for large datasets and is costly for NNs.
Instead, we simply set ˆq = 0 , and optimise the ﬁrst quantile once only.
We show this is roughly equivalent in Appendix A.2.
Unfortunately, this adaption is still inefﬁcient.
Inspecting Algorithm 1, one notes that M NNs must be trained sequentially, meaning time complexity (both at training and at test) and memory complexity is
O ( M ) .
This is not the case when doing quantile regression with NNs without censoring – a single NN typically uses multiple outputs, one per quantile, and their losses can be minimised simultaneously [Taylor, 2000].
For modest numbers of quantiles, this results in little overhead relative to a single
NN, since the extra quantile outputs require only linear heads, while the bulk of the parameter count and computation time arise from the shared trunk of the NN.
Note that these inefﬁciencies are less problematic for linear models – optimisation of Eq.
4 can be performed quite rapidly, and memory complexity cannot obviously be improved beyond O ( M ) .
As such, it is the use of NNs that motivates the search for an alternative algorithm.
3.2 CQRNN Algorithm
Our proposed algorithm uses a single NN outputting a grid of quantile estimates that can be optimised simultaneously.
This results in a more efﬁcient algorithm in terms of training time, test time, and memory.
Concretely, if the NN is a single hidden-layer NN with H neurons, then the model is deﬁned, ⌧ ( x j , ✓ ) =
✓ | 1 , ⌧  ( [-θ (cid:124) 
-]✓ | 0 x i ) , with ✓ 0 2 R D ⇥ H shared for all estimates and ✓ 1 , ⌧ 2 R H ⇥ 1 for the speciﬁc output head, and non-linearity,  .
The key insight behind our algorithm is that estimates of the censored weights, ˆw , can be bootstrapped from the model while it learns, without requiring sequential optimisation.
Intuitively, as the model trains, the estimates of the censored quantiles, ˆq , and hence the censored weights, ˆw , improve, and these feedback to create a more accurate loss function, allowing further improvement of the model.
We summarise the method in Algorithm 2.
The quantile of a censored datapoint is estimated as whichever NN output is closest under the current model parameters.
Whilst the algorithm is written in two steps, getting up-to-date estimates of censored quantiles, ˆ q j , requires only a forward pass, which is done at the point of loss optimisation anyway, so actually comes ‘for free’.
In our implementations, the algorithm is extended to leverage mini-batches and modern optimisers (e.g. Adam).
Analysis of CQRNN
The previous section introduced the CQRNN algorithm, motivating it intuitively, while this section considers its correctness.
Providing precise analytical guarantees is very challenging, both due to the involvement of non-convex NNs, and the bootstrapping nature of the algorithm.
We therefore consider more holistic ways of understanding and providing conﬁdence in the method.
We ﬁrstly show that it can be interpreted as a ﬂavour of EM (‘generalised hard’ EM) – connecting to this well-studied class of algorithms provides some justiﬁcation of CQRNN’s design.
We secondly validate the way the model bootstraps its own censored quantile estimates by describing a ‘self-correcting’ property.
Algorithm 2 CQRNN algorithm.
Require: Dataset D , one parametric model ( · ) with randomly initialised parameters ✓ that can output M quantile predictions, quantiles to be estimated grid ⌧ , learning rate ↵ , pseudo y value y ⇤ .
S censored { i 2 { 0 , 1 , . . .
, N } :  i = 0 } , S observed { i 2 { 0 , 1 , . . .
, N } :  i = 1 } while not convereged do 1.
Hard expectation step for j 2 S censored do ˆ q j
[Equation]
Estimate quantile of censored data under current model
[Equation]
Partial maximisation step
[Equation]
. Gradient step to minimise Eq.4.1 Interpretation as Expectation-Maximisation
EM is a two-stage iterative optimisation technique for ﬁnding maximum likelihood solutions.
It can be useful when one has both observed and latent variables, and where knowing the latent variables would allow straightforward optimisation of a model’s parameters and vice versa.
In our problem, we interpret the weights of the censored datapoints, w , as latent variables, observed variables as x & y , and the NN parameters ✓ are to be optimised.
Under our notation, the vanilla EM procedure [Bishop, 2006] can be written as follows (using a randomly initialising ✓ old ),
Expectation step.
Evaluate, p ( w | x , y, ✓ old ) , to set up the expectation,
[Equation]
Maximisation step.
Perform optimisation of the expected likelihood.
[Equation]
Set ✓ old ✓ new and repeat until convergence.
We now show how our CQRNN method in Algorithm 2 can be interpreted as a ﬂavour of EM, speciﬁcally generalised hard EM.
The key steps in this interpretation are: 1)
The likelihood for datapoints at each quantile are chosen to follow an asymmetric Laplace distribution.
The expectation step is treated as a hard assignment under the current NN parameters.
The likelihood is only partially maximised at each iteration.
Likelihood Form
A likelihood function is required for both the E & M steps.
Theorem 1 shows how Eq.
4 can be interpreted as a negative log likelihood.
Theorem 1.
Let the likelihood for each datapoint at each quantile be an asymmetric Laplace distribution with scale,  = p ⌧  ⌧ 2 , and asymmetry, k = ⌧ / p ⌧  ⌧ 2 .
The negative log likelihood is,
[Equation]
Proof sketch.
The asymmetric Laplace distribution has been used in Bayesian treatments of quantile regression without censoring [Yu and Moyeed, 2001].
We extend this to Portnoy’s loss with a weighted likelihood form using the censored weights.
Appendix A provides the full proof.
Expectation as a Hard Assignment of Latent Variables
In some models, such as a Gaussian mixture model, the expectation step can be computed through analytical evaluation of the latent posterior distribution [Bishop, 2006].
In our case, p ( w | x , y, ✓ old ) , is intractable, though Eq. 4 depends linearly on each element w j , so we need only consider the expectation, E w j |
x ,y, ✓ old [ w j ] .
The algorithm can then be interpreted as making an assignment under the current model parameters,
[Equation]
This is a ‘hard’ assignment, where latents are assigned to the most likely values under the current model (e.g. datapoints are attributed to the nearest clusters in the K-means algorithm
[Bishop, 2006]), and gives rise to a class of algorithms termed hard EM [Samdani et al., 2012].
Note that the in Eq.
6 we can rewrite, log p ( y, w | x , ✓ )
= log p ( w ) p ( y | x , ✓ , w ) .
Since we choose only a single setting for w in the outer expectation, this dependence disappears, and we require maximisation of, E w | x ,y, ✓ old [log p ( y | x , ✓ , w )] .
Partial Maximisation
A second departure from the standard EM algorithm, is that Algorithm 2 performs only a partial maximisation of the likelihood, taking a single gradient step.
This has been shown to provide similar guarantees, and is termed ‘generalised’ EM [Neal and Hinton, 1998].
One could consider an alternative version of the algorithm that is closer to the standard EM procedure, where the estimated quantiles of censored data points are ﬁxed while the maximisation step of NN parameters is run to convergence over multiple training epochs, before the estimates are updated in the expectation step, and repeating.
In our case this is less effective – the maximisation requires running a forward pass through the NN, and obtaining estimated quantiles under current NN parameters once this is done is trivial, so we can access up-to-date estimates essentially for free.
Appendix
Figure 5 empirically demonstrates that CQRNN’s partial maximisation approach produces the fastest convergence.
4.2 Self-Correcting Property
The CQRNN bootstraps weight estimates from the current model as it trains.
It’s not immediately obvious why this bootstrapping approach should converge to something sensible – what if bad initial weight estimates lead to worse ones?
As a second insight into the CQRNN algorithm, Theoremshows that when a censored weight, ˆ w j , is estimated incorrectly, the algorithm acts in a way to adjust the estimated quantiles in a favourable way – we refer to this as ‘self-correcting’.
Theorem 2.
If ˆ q j is underestimated, one iteration of the algorithm acts to increase the quantile predictions, ˆ y j, ⌧ , by the same amount, or even higher, than if the weight had been correct.
If ˆ q j is overestimated, ˆ y j, ⌧ , is decreased by the same amount, or even lower, than with the correct weight.
Proof sketch.
Denote q
j the true quantile that censored datapoint j is censored in, and w j the corresponding true weight.
We derive the expression for the gradient wrt the predicted quantiles, ˆ y j, ⌧ , ﬁnding that if ˆ w j is underestimated it holds that, @ L Port.
( ✓ , D , ⌧ , ˆ w ,y ⇤ ) @
ˆ y j, ⌧  @ L Port.
( ✓ , D , ⌧ , w ,y ⇤ ) @
ˆ y j, ⌧ , and hence gradient descent applies an adjustment in the desired direction of equal or greater magnitude than if the true weight had been used.
The reverse holds for overestimated ˆ w j .
Appendix A provides a full proof.
Related Work
Survival analysis and NNs.
Following the widespread success of deep NNs over the past decade, there has been a wave of research applying NNs to survival analysis – for instance by modifying the CoxPH model [Katzman et al., 2018], or framing the task as ordinal classiﬁcation [Lee et al., 2018].
Closer to our work are methods that use NNs to output parameters of a distribution such as the Weibull [Martinsson, 2016], seek robust training objectives for these models [Avati et al., 2019], or help with their optimisation [Tang et al., 2022].
Our work stands out as offering a way to directly estimate the target variable at pre-speciﬁed quantiles, without enforcing any distributional assumption.
A limitation is that the distribution may only be predicted at these quantiles (unless additional assumptions are made to allow interpolation between these).
Quantile regression and NNs.
Quantile regression has proven an attractive option to enable NNs to move beyond point predictions.
This allows quantiﬁcation of a NN’s aleatoric uncertainty [Tagasovska and Lopez-Paz, 2019].
It is attractive due to its straightforward implementation, and avoidance of any distributional assumption.
For example, it has found use in reinforcement learning to capture the distribution of rewards, rather than just the mean [Dabney et al., 2017].
Since NNs are ﬂexible function approximators, particular attention has been paid to the crossing quantile problem [Bondell et al., 2010, Zhou et al., 2020, Brando et al., 2022].
CQRNN borrows ideas from this line of work, and further tackles the challenge of learning quantiles under censored data.
Censored quantile linear regression.
There is much work on censored quantile regression methods for linear models.
Powell [1986] developed an estimator under ﬁxed-value censoring which can be implemented with an algorithm from Fitzenberger [1997].
Portnoy [2003] developed an estimator under random censoring based on the KM estimator, while Peng and Huang [2008] developed an alternative based on the Nelson–Aalen (NA) estimator.
KM and NA are closely related, and Portnoy and Peng’s methods have been reported to offer similar empirical performance [Koenker, 2008].
Koenker [2022] provides all above methods in the popular ‘ quantreg ’ R package.
Other notable methods include Yang et al.
[2018], based on the data augmentation algorithm, and Wang and Wang [2009], whose estimator is similar to Portnoy’s
but utilises local estimates of the KM, computed with a kernel method.
See Peng [2021] for a review of the broader area.
Our work allows modelling of ﬂexible non-linear quantile functions, leveraging the powerful representation learning abilities of
NNs.
Although, this sacriﬁces the interpretability of coefﬁcients of linear models.
Censored quantile regression and NNs.
Only a small amount of work has been done in this area.
The ‘ qrnn ’ R package [Cannon, 2019] offers the ability to train NNs under ﬁxed-value left censoring, adopting an idea from the linear setting [Friederichs and Hense, 2007, Cannon, 2011]. Huttel et al.
[2022] explore our objective but assume censoring times, c i , are available for both censored and uncensored data points (i.e. the censoring distribution is known).
We do not require this assumption.
DeepQuantReg [Jia and Jeong, 2022] tackles the same objective as this paper.
They showed that improvements in quantile estimation can be obtained relative to naive methods.
Their work differs from ours signiﬁcantly – they base their method around an estimator for the median from Huang et al.
[2007], requiring an assumption that the censoring distribution is independent of covariates.
Experiments
This section empirically investigates several questions.
Q1)
How does the proposed CQRNN method compare with existing methods?
This is done qualitatively on synthetic 1D functions in Section 6. and quantitatively on synthetic and real datasets in Section 6.2.
Q2)
How does the sequential grid algorithm compare to the CQRNN algorithm, both in terms of predictive accuracy and efﬁciency?
Explored in Section 6.3 Q3)
How is the CQRNN algorithm affected by its hyperparameters?
We investigate the impact of grid ﬁdelity and y ⇤ in Section 6.4.
All experiments use fully-connected NNs with two hidden layers of 100 neurons, except for SurvMNIST, when three convolutional layers are used.
Grid size M is set to 5, 9 or 19 depending on dataset size.
Appendix B contains further details on hyperparameter settings, metrics, and datasets.
Appendix
C presents further results.
Datasets.
We use three types of dataset.
Type 1) Synthetic target data with synthetic censoring.
Type 2)
Real target data with synthetic censoring.
Type 3)
Real target data with real censoring.
Whilst type 3 captures the kind of datasets we care about most, evaluation of quantiles is challenging since in real-world survival data the target conditional quantiles are not obtainable even at test time [Li and Peng, 2017].
In contrast, type 1 offers access to the ground truth quantiles for clean evaluation, but properties of these datasets may be less realistic.
We introduce type 2 as a middle ground, which takes a real-world dataset without censoring, and synthetically overlays a censoring distribution to create training data.
At test time, we have the option of not applying the censoring, providing samples from the clean target distribution, allowing clearer evaluation.
We summarise all datasets in Table 2.
Full details are in Appendix B.4.
Our experiments exceed the number and variety of datasets used in popular recent works [Goldstein et al., 2020, Zhong et al., 2021]
[Table]
Metrics.
Our objective is to measure how closely a model’s predicted quantiles match those of the ground truth target distribution.
We favour different metrics for each dataset type.
For type 1 datasets, since targets are generated synthetically it is possible to analytically compute the ground truth target quantile for an input x i , which we denote y i, ⌧ .
We compute the mean squared error (MSE) between the predictions and the ground truths across three quantiles, ⌧ 2 [0 .
9] , (we ensure these are always present in grid ⌧ ).
This is our ﬁrst-choice metric when available.
True quantile MSE (TQMSE) :
= 1 N
[Equation]
In type 2 datsets samples from the uncensored target distribution are available but not the synthetic generating function.
Our preferred metric is the checkmark loss across the three quantiles.
Uncensored quantile loss
[Equation]
For datasets of type 3, we use two metrics.
The concordance index (C-index) is computed using the median ( ⌧ = 0 . 5 ).
But this may not reveal anything about systematic bias of different models, nor about other quantiles ⌧ 6 = 0 .
We secondly use censored D-Calibration (CensDCal)
[Haider et al., 2020] which measures whether the empirical proportion of datapoints falling between pairs of consecutive quantiles, matches the target proportion, ⌧ j +1  ⌧ j .
MSE of the deviation is then computed.
There is also an uncensored version (UnDCal), which we can compute for type 1 &datasets.
Appendix B.2 gives further details.
Baselines.
We compare against three methods that can be used to predict the quantiles of a target distribution using a NN.
Excl.
Censor – a method that naively excludes censored datapoints from the training data, optimising the loss in Eq.
DeepQuantReg – the only existing method in the literature proposing explicit output of quantiles from a NN on censored data [Jia and Jeong, 2022].
LogNorm MLE – A NN outputting parameters of a lognormal distribution, that’s trained via maximum likelihood estimation (MLE) (details in Appendix B.3), and is a standard baseline in related work (e.g. [Avati et al., 2019, Goldstein et al., 2020]) since this distribution often suits properties of real-world time-to-event survival data [Kleinbaum and Klein, 2012].
6.1 Qualitative 1D Analysis
Figure 1 visualises the quantiles predicted by CQRNN and baseline methods for 1D datasets (Table 2 describes dataset functions, Figure 3 visualises further 1D datasets).
Each dataset containsdatapoints drawn from, x ⇠ U (0 , 2) and grid ⌧ 2 { 0 . 1 , 0 . 3 , 0 . 5 , 0 . 7 , 0 . 9 } .
CQRNN usually learns quantiles that are closer to the ground truth than baseline methods, particularly at higher quantiles.
Excl.
censor systemically makes underpredictions, and this worsens at higher quantiles since larger values of y are more likely to be censored and excluded.
DeepQuantReg avoids this systemic underprediction, but appears to introduce bias of its own.
LogNorm MLE provides variable results – on Norm uniform it fails with excessive variance (an issue also observed by Avati et al.
[2019]), but on LogNorm, since the target distribution matches the distribution output by the NN, it performs well.
6.2 Benchmarking
[Figure]
( l o w e r
[Figure]
tt
[Figure]
Our main experiment benchmarks CQRNN against all baselines across a wide variety of datasets, covering various domains, sizes, dimensionalities and censoring proportions (see
Table 2).
Some hyperparameter tuning was performed for each method (Appendix B.1).
Figure 2 plots the preferred metric for each dataset type, though patterns were consistent across metrics – Table 4 provides a full breakdown.
For type 1 (synthetic) datasets, the quantitative results follow our qualitative observations, with
CQRNN producing lowest TQMSE on all but the datasets generated from a Log Normal distribution, when it sometimes placed second behind
LogNorm MLE.
CQRNN also produces the lowest UQL on all datasets of type 2.
LogNorm MLE performs poorly on these, presumably since they are not typical time-to-event datasets which are well captured by the Log Normal distribution.
For type 3 datasets, CQRNN and LogNorm
MLE usually perform best in terms of C-index, with error bars tending to overlap.
In CensDCal,
CQRNN consistently perform best, matched only by LogNorm MLE on two datasets, when error bars overlap.
Excl.
censor is the weakest method, with DeepQuantReg midway between it and CQRNN.
6.3 Comparison of Sequential Grid and CQRNN
We ran a head-to-head comparison of CQRNN and the sequential grid algorithm on all type 1 synthetic datasets, and all type 3 real datasets.
Table 1 compares on type 1 datasets, along with statistical tests of signiﬁcance for TQMSE differences (200 random seeds).
Appendix Table 5 provides results for the type 3 datsets, reporting C-index and CensDCal (50 random seeds).
Signiﬁcance tests are described in Appendix B.1.
CQRNN delivers beneﬁts in speed up both at training and test time, running an order of magnitude faster than the sequential grid algorithm.
CQRNN also dramatically reduces model size.
The magnitude of these beneﬁts is largely determined by the size of the quantile grid, M , which is explicit in the sequential algorithm’s time and space complexity, O ( M ) , but largely avoided in CQRNN.
Differences in the quality of quantiles of the two algorithms is usually slight (after-all both leverage the same estimator), though CQRNN shows statistically signiﬁcant gains on 8/14 type 1 datasets in terms of TQMSE, 4/7 type 3 datsets in terms of C-Index, and 7/7 type 3 datasets in terms of CensDCal.
We hypothesise that having a single NN output all quantiles provides a helpful inductive bias, encouraging similarity between adjacent quantiles, that’s not present in independently trained NNs.
6.4 Hyperparameter Investigation
The CQRNN algorithm includes two hyperparameters – the grid of quantiles, grid ⌧ , and the pseudo y value, y ⇤ .
We ran ablations to empirically investigate the effect of these, as well as remedies for the ‘crossing-quantile’ problem, and a comparison of partial vs. full optimisation.
Here we summarise our ﬁndings, Appendix C.1 provides full details.
We tested grid sizes, M 2 { 9 , 19 , 39 } , on several type 1 datasets.
In general, a ﬁner grid (larger M ) is slightly beneﬁcial, though there is variance between datasets and gains are sometimes only seen for larger datasets (>5,000 datapoints).
Selection of y ⇤ requires some care for CQRNN.
We deﬁned it in terms of the maximum y value in the training set, y ⇤ =
c y ⇤ max i
y i , for a hyperparemeter, c y ⇤ > 1 .
Performance can be improved by tuning c y ⇤ for each dataset, but using c y ⇤ = 1 . 2 provided consistently reasonable results – this was the value used in our benchmarking experiments.
We trialled two methods for combating the crossing-quantile problem.
1) Adding a crossing loss penalty.
2) Constraining the NN architecture to output quantiles adding to the previous prediction.
Neither of these methods signiﬁcantly affected performance.
Discussion & Conclusion
This paper has taken a popular idea from survival analysis, Portnoy’s censored quantile regression estimator, and shown how it can be efﬁciently combined with NNs in a new algorithm, CQRNN.
We provided theoretical insight by interpreting it as a ﬂavour of EM.
Empirically the method outperformed existing approaches, consistently producing more accurate quantile estimates across a range of synthetic and real datasets.
For example, across datasets of type 2 and 3, CQRNN was best calibrated in 10 out of 12 instances (see CensDCal in Table 4).
Limitations.
Firstly, our theoretical results have not said that solutions will converge on a global optimum.
Secondly, we have only tested CQRNN on a modest number of real-world datasets, many drawn from the biomedical domain.
It’s possible that some datasets may cause CQRNN problems, as we found on BreastMSK (Figure 2).
In particular, cases where higher quantiles are undeﬁned should be handled with care.
Conclusion.
To summarise, our work contributes toward unlocking the beneﬁt that modern machine learning could bring to important domains such as healthcare and machinery prognostics.
By outputting quantiles, CQRNN naturally communicates a measure of uncertainty in its predictions, which makes it particularly suitable to these high-stakes applications, and a valuable addition to the toolkit combining deep NNs with survival analysis.
[Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure]
[Table] [Figure]
U (0 , 3 . 5) LogNorm same   0.
Equal to target dist.
Type 2 – Real datasets with synthetic censoring Housing 8 0.
Real
U (0 ,c 1 ) Protein 9 0.
Real
U (0 ,c 2 ) Wine   0.
Real
U (0 ,c 3 ) PHM 21 0.
Real
U (0 ,c 4 ) SurvMNIST 28 ⇥ 28 480000.
One Gamma
[Table]
per MNIST class
U (0 ,c 5 ) Type 3 –
[Table]
with
[Table]
METABRIC   0.
Real
Real WHAS   0.
Real
Real SUPPORT   0.
Real
Real GBSG   0.
Real
Real TMBImmuno   0.
Real
Real BreastMSK   0.
Real
Real LGGGBM   0.
Real Real
Acknowledgement
This work was supported by NSF of China Projects (Nos. 62061136001, U19B2034, U1811461,
U19A2081); Beijing NSF Project (No. JQ19016); a grant from Tsinghua Institute for Guo Qiang; and the High Performance Computing Center, Tsinghua University.
J.Z was also supported by the
XPlorer Prize.
[-Epochs were set via the formula epochs = 
-]