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 ∈ R D is the input, y i ∈ 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)
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 ∀
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, τ ∈ (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,τ = θ (cid:124)
τ x i , with, θ τ ∈ 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]
= (cid:88)
i ∈S observed
[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 ∗ (cid:29) 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 ∈ R N c and q ∈ 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].
[Equation]
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 ∈ { 0 , 1 , . . .
, N } : ∆ i = 0 } , S [Table]
← { i ∈ { 0 , 1 , . . . , N } : ∆ i = 1 } , K cross ← ∅ for i = 0 to M
[Table]
1 do
[Table]
if i = [-
-][Table]
then (cid:46) Initialise quantile estimates to zero [Table] else (cid:46) Find indices which have been crossed and update their estimated quantiles
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] (cid:46) [Table] ˆw = 0 for uncrossed datapoints
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] (cid:46) 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 .
CQRNN (ours)
[Figure]
censor DeepQuantReg LogNorm MLE 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
CQRNN recovers quantiles closest to the ground truth on most datasets.
Censored Quantile Regression and Neural Networks
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 , θ )
= θ (cid:124) 1 ,τ φ ( θ (cid:124)
x i ) , with θ 0 ∈ R D × H shared for all estimates and θ 1 ,τ ∈ 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 ∈ { 0 , 1 , . . .
, N } : ∆ i = 0 } , S observed ← { i ∈ { 0 , 1 , . . .
, N } : ∆ i = 1 } while not convereged do Hard expectation step for j ∈ S censored do ˆ q j [Equation] (cid:46) Estimate quantile of censored data under current model
[Equation]
Partial maximisation step
[Equation]
(cid:46) 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, λ = √ τ − τ 2 , and asymmetry, k = τ/ √ τ − τ 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].
DeepQuantReg [Jia and Jeong, 2022] is the only prior work taking on 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.
Code is available at https://github.
com/TeaPearce/Censored_Quantile_Regression_NN .
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, τ ∈ [0 . 1 , 0 . 9] , (we ensure these are always present in grid τ ).
This is our ﬁrst-choice metric when available.
True quantile MSE (TQMSE) :
= 1 N (cid:80) τ ∈ [ [Equation]
] (cid:80)
[-=
-][Equation] (ˆ [Equation]
− y [Equation] τ ) 2 .
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 (UQL) :
= 1 N (cid:80) τ ∈ [0 . 1 , 0 . 5 , 0 . 9] (cid:80) Ni =
1 ρ τ ( y i , ˆ y i,τ ) .
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 τ (cid:54)
= 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 τ ∈ { 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 ∈ { 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
[Table]
[Figure]
[Table]
LogNorm same   0.
Equal to target [Table] [Figure] [Table]
with
[Table]
Housing 8 0.
Real U [Table] [Figure]
[Table]
1 ) Protein 9 0.
Real
U [Table] [Figure] c 2 )
Wine   0.
Real
U (0 ,c 3 ) PHM 21 0.
Real U ( [Figure] [Table]
4 ) SurvMNIST
[Figure]
One
[Table]
[Figure] [Table] class [Table]
) Type
[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
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. A Analytical Results A.1 Proofs Theorem 1. Let the likelihood for each datapoint at each quantile be an asymmetric Laplace distribution with scale, λ = √ τ − τ 2 , and asymmetry, k = τ/ √ τ − τ 2 .
The negative log likelihood is,
[Equation]
Proof.
Deﬁne the likelihood over all quantiles of interest, and split censored datapoints into two pseudo datapoints, one at the censoring location and one at the large pseudo value y ∗ , to give a weighted likelihood,
[Equation]
We write the asymmetric Laplace density with ˆ y j,τ as the location parameter and scale λ and asymmetry k ,
[Equation]
else
[Equation]
[Equation]
Setting λ = √ τ − τ 2 and k = τ/ √ τ − τ 2 and rearranging,
[Equation]
[Equation]
Taking the logarithm of Eq.
13 we have,
[Equation]
Eq.
16 may be substituted into this if all likelihoods of observed and censored pseudo datapoints are chosen to follow asymmetric Laplace distributions.
Taking the negative then recovers the theorem’s result.
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.
Remark.
Note that if a quantile ˆ q j is underestimated, it’s desirable to increase the quantiles relating to that datapoint, for input x j and predictions ˆ y j,τ .
If a quantile ˆ q j is instead overestimated, it’s desirable to decrease the quantiles for input x j and predictions ˆ y j,τ .
We refer to this desired behaviour as ‘self-correcting’.
Proof.
Denote q j the true quantile that censored datapoint j is censored in, and w j = ( τ − q j ) / (1 − q j ) the corresponding true weight.
We now consider one iteration of the algorithm in the case that the estimated quantile is underestimated.
Model underpredicts the censored quantile, ˆ q j = q j − (cid:15) , for some (cid:15) > 0 .
The corresponding weight is also underestimated, ˆ q j < q j = ⇒
ˆ w j < w j , shown in lemma 2.1.
Lemma 2.3 shows that, for censored datapoint j , the gradient of Eq.
4 wrt the quantile prediction ˆ y j,τ is,
[Equation]
Hence, if ˆ w j is underestimated it holds that, ∂ L Port.
( θ, D ,τ
[Equation]
ˆ
[Equation]
∂
[Equation]
When this gradient is used for optimisation, this has the effect of increasing the quantile prediction, ˆ y j,τ , by either the same amount, or even higher, than if the weight had been correct.
Similar (reversed) logic applies if the quantile is initially overestimated, ˆ q j = q j + (cid:15) , (e.g. lemma 2.2) which encourages decreasing the quantile predictions, ˆ y j,τ , by the same amount or lower than with correct weights.
Hence, weight estimates will be improved in future iterations of the algorithm.
Note that this applies to all quantiles, τ ∈ grid τ .
Lemma 2.1.
Let, ˆ q j = q j − (cid:15) , and (cid:15) > 0 .
It holds that, ˆ q j < q j = ⇒ ˆ w j < w j , for, τ ∈ (0 , 1) , and, ˆ q j , q j ∈ (0 , 1) .
Proof.
From the deﬁnition of the weights, ˆ w j < w j = ⇒ ( τ − q j − (cid:15) ) /
(1 − q j − (cid:15) ) <
( τ − q j ) / (1 − q j ) .
Let a :=
τ − q j and b := 1 − q j .
Note that, a < b , since by assumption, τ < 1 .
We must show
[Equation]
Lemma 2.2.
Let, ˆ q j = q j + (cid:15) , and (cid:15) > 0 .
It holds that, ˆ q j > q j = ⇒ ˆ w j > w j , for, τ ∈ (0 , 1) , and, ˆ q j , q j ∈ (0 , 1) .
Proof.
This proof follows lemma 2.1.
We now have,
[Equation]
which holds since, a < b.
[Equation]
[Equation]
Lemma 2.3.
The partial derivative of Portnoy’s loss wrt the predicted quantile, for censored datapoint j , is given by,
[Equation]
Proof.
Recalling that,
[Equation]
24) we have,
[Equation]
We are interested in the derivative for the censored portion of Portnoy’s loss in Eq. 4,
[Equation]
[Equation]
Note that we always choose, y ∗ > y j .
Hence ˆ y j,τ < y j = ⇒
ˆ y j,τ < y ∗ , so there are three cases to consider.
Case 1, ˆ y j,τ < y j , case 2, y j ≤ ˆ y j,τ < y ∗ , case 3 y ∗ ≤ ˆ y j,τ .
For case 1,
[Equation]
For case 2,
[Equation]
For case 3,
[Equation]
A.2 First Iteration of the Sequential Grid Algorithm
In this section we compare the procedure proposed by Portnoy [2003] to ﬁnd the ﬁrst quantile predicted at, τ 0 := grid τ [0] , with the procedure we propose in the sequential grid algorithm for NNs (Algorithm 1).
We show that these two procedures produce equivalent gradients.
Portnoy’s procedure.
Portnoy [2003] require that no censored datapoints lie below the ﬁrst quantile, and propose deleting these from the dataset when this does occur, as follows.
Set w = 1 for all censored datapoints in dataset. 2.
Optimise L Port.
(Eq. 4) for τ 0 .
Find all censored datapoints below τ 0 , B ← { j ∈ S censored :
y j < ˆ y j,τ 0 } .
If B is empty then exit.
Exclude all elements of B from the dataset and repeat.
This optimisation could be repeated many times.
We’d like to avoid this since training NNs on potentially large datasets can be costly.
Sequential grid for NNs procedure.
Algorithm 1 instead simply sets q = 0 , and optimises the ﬁrst quantile once only.
Set q = 0 for all censored datapoints in dataset.
Optimise L Port.
(Eq. 4) for τ 0 .
Justiﬁcation.
We now justify why this is a reasonable approximation.
First note that when ˆ q j =we have, ˆ w i = τ
− ˆ q i 1 − ˆ q i = τ .
Using lemma 2.3 we can compare the gradients for each procedure.
For Portoy’s procedure, when ˆ w j = 1 ,
[Equation]
next iteration
[Equation]
For Algorithm 1, when ˆ q j = 0 ,
[Equation]
At ﬁrst look, the gradients appear to differ in the case y j ≤ ˆ y j,τ < y ∗ .
But when a datapoint triggers this criteria in Portnoy’s procedure, it will be excluded and the model retrained its
B Experimental Details
This section provides further details about all experiments run.
Our code base uses the PyTorch framework.
Hyperparameters in Appendix B.1.
Metrics in Appendix B.2.
Baselines in Appendix B.3.
Datasets in Appendix B.4.
Hardware.
We used an internal cluster for experiments, utilising machines with four GPUs and[-
-]
CPU cores.
Most of our datasets used fully-connected NNs, which were trained on CPU, while GPUs were used for the SurvMNIST experiments.
B.1 Full Hyperparameter Details
Below we list hyperparameter settings used and where applicable the tuning protocols followed.
B.1.1 Qualitative 1D Analysis
Section 6.1 experiment.
All methods used the same optimisation procedure and NN architecture, without tuning.
• Training dataset size: 500, where, x ∼ U (0 , 2) • Epochs:• Optimiser: Adam • Learning rate: 0.01 (decreased 70% and 90% of the way through training) •
Batch size:• Weight decay: 0.
• NN architecture: Fully-connected, two hidden layers of 100 hidden nodes, GeLU activations • y ∗ = 1 . 2 × max i
y
i • grid τ ∈ { 0 . 1 , 0 . 3 , 0 .
B.1.2 Benchmarking
Section 6.2 experiment.
Experiments were repeated over 10 random seeds.
Hyperparameter settings.
• Training dataset size: various – see Table• Test dataset size: various – see Table• Epochs: ∈ { 10 , 20 , 50 , 100 } • Optimiser: Adam • Learning rate: 0.01 for fc NN, 0.001 for CNN (decreased 70% and 90% of the way through training) •
Batch size:• Weight decay: 0.
• Default NN architecture: Fully-connected, two hidden layers of 100 hidden nodes, ReLU activations • CNN architecture for SurvMNIST following Goldstein et al.
[2020]: Conv2D[64, (5 × 5)]
→ ReLU → Dropout(0.2)
→ AveragePool(2 × 2) → Conv2D[128, (5 × 5)]
→ ReLU → Dropout(0.2) → AveragePool(2 × 2) → Conv2D [Table] 256, (2 × 2)]
→ ReLU → Linear • y ∗ = 1 . 2 × max i y i • Grid size M ∈ { 6 , 10 , 20 } • Dropout ∈ { True , False }
Tuning process for epochs and dropout:
• For the real type 2 and type 3 datasets, we tuned number of epochs ∈ { 10 , 20 , 50 , 100 } and dropout ∈ { True , False } for each dataset for each method.
We used three random splits as a validation (but not overlapping with the random seeds used in the ﬁnal test run).
• Epochs was ﬁxed to 100 and dropout disabled for: Norm linear, Norm non-linear, Exponential, Weibull, LogNorm, Norm uniform.
• Epochs was ﬁxed to 20 and dropout disabled for: Norm heavy, Norm medium, Norm light, Norm same.
[Table]
• Epochs was ﬁxed to 10 and dropout disabled for: LogNorm heavy, LogNorm medium, LogNorm light, LogNorm same.
We ﬁxed the grid size according to an estimate of how densely the datapoints covered the input space (a rough consideration of dataset size and number of features):
• Grid size M = 5 for smaller datasets with more features: WHAS, SUPPORT, GBSG, TMBImmuno, BreastMSK, LGGGBM, METABRIC.
• Grid size M = 9 for medium datasets or smaller datasets with less features: Norm linear, Norm non-linear, Exponential, Weibull, LogNorm, Norm uniform.
• Grid size M = 19 for larger datasets or those with less features: Norm heavy, Norm medium, Norm light, Norm same, LogNorm heavy, LogNorm medium, LogNorm light, LogNorm same, Housing, Protein, Wine, PHM, SurvMNIST.
B.1.3 Comparison of Sequential Grid and CQRNN
Section 6.3 experiment.
This was carried out under the same protocol as for the main benchmarking, but repeated over a larger number of seeds (200 for type 1 datasets, 50 for type 3 datasets).
To obtain 95% conﬁdence intervals, we compute the standard error of the difference between means, and multiply it by a two-sided t-statistic, with (number of seeds − 1 ) degrees of freedom at the α = 0 .signiﬁcance level. If zero falls within this conﬁdence interval, the difference is deemed not signiﬁcant.
B.1.4 Hyperparameter Investigation
Section 6.4 experiment.
Hyperparameters are as for the main benchmarking except M and y ∗ were varied as stated in the text.
Epochs were set via the formula epochs =
× 200 /N ensuring the same number of gradient updates were made on each run.
Experiments were repeated overrandom seeds for the grid investigation, and ten random seeds for the y ∗ investigation.
B.2 Metrics
This section brieﬂy expands upon the metrics introduced in Section 6.
Table 3 summarises the availability on each metric for each dataset type.
The computation of each is also detailed below.
True quantile
[Equation]
Uncensored quantile loss (UQL
[Equation]
Uncensored D-Calibration (UnDCal) :
= 100 ×
[Equation]
Censored D-Calibration (CensDCal) :
= 100 ×
[Equation]
where, Goldstein et al.
[2020] deﬁnes, ξ = (cid:88)
i ∈S observed
I [ˆ y i,τ j < y i ≤ ˆ y i,τ j +1 ]+ (cid:88) i ∈S censored ( τ j +1 − q i ) I [ˆ y i,τ j < y i ≤ ˆ y i,τ j +1 ] 1 − q i + ( τ j +1 − τ j ) I [ q i < τ j ] 1 − q i . We increase the magnitude of DCal metrics by 100 × to make the numbers of similar order to TQMSE and UQL.
B.3 Baselines
This section provides some extra detail about the LogNorm MLE baseline.
For this method, we use a NN with two outputs, representing the mean, µ , and standard deviation, σ , of a Log Normal distribution, i.e. log y ∼ N (ˆ µ, ˆ σ 2 ) .
We pass the output representing the standard deviation prediction through a SoftPlus to ensure it is always positive and differentiable.
The maximum likelihood estimation loss is then,
[Equation]
where the likelihood and CDF follow the analytical expressions for the Log Normal distribution.
At evaluation time, we use the SciPy package to compute the quantiles from the predicted Log Normal distribution that correspond to those in grid τ .
This allows a like-for-like comparison with our other baselines.
B.4 Dataset Details
This section provides further detail about the source of each dataset used.
All real-world datasets were taken from open-access repositories, and had already been anonymised.
For real-world datasets, we do not follow any previous test/train splits, rather we randomly shufﬂe the data for each run, selecting 80% for training and 20% for testing.
B.4.1 Type 1 Datasets
details the generating target and censoring distributions used, as well as numbers of test and train datapoints.
Inputs were always generated uniformly via, x ∼ U (0 , 2) D for D features.
B.4.2 Type 2 Datasets
Table 2 details dataset sizes and number of features.
All datasets used a uniform censoring distribution, c i ∼ U (0 , c ) , where c was selected to be equal to the 90th percentile of the target distribution for SurvMNIST, and c = 1 . 5 max
i y i for the other type 2 datasets.
Housing was sourced from
Scikit Learn datasets, while Protein, Wine and PHM were sourced through OpenML https: //www.openml.org/ .
• Housing Target is median house prices.
Retrieved from https://scikit-learn.org/ stable/datasets/real_world.html#california-housing-dataset .
• Protein Target is RMSD.
OpenML lookup ID is physicochemical-protein .
• Wine Target is quality of wine.
OpenML lookup ID is wine quality .
• PHM (prognostics health management) target is breakdown time of simulated machines.
OpenML lookup ID is NASA PHM2008 .
Our ﬁnal type 2 dataset is slightly different, since we don’t use the provided labels directly.
• SurvMNIST appeared in Goldstein et al.
[2020], who adapted it from Sebastian Pölsterl’s blog:
https://k-d-w.org/blog/2019/07/ survival-analysis-for-deep-learning/ .
It uses the standard MNIST dataset http://yann.lecun.com/exdb/mnist/ , but targets are drawn from a Gamma distribution, with different parameters per class.
Goldstein et al.
[2020] used a small variance ﬁxed across classes, with means
∈
For
our purposes, we are most interested in how well methods capture variance, so we vary it, variance ∈ [0 . 1 , 0 .
B.4.3 Type 3 Datasets
We provide a brief overview of each dataset.
Four datasets – GBSG, METABRIC, SUPPORT, WHAS – were all retrieved from https://github.com/jaredleekatzman/DeepSurv/tree/master/ experiments/data .
Katzman et al.
[2018] provides a detailed introduction to these datasets.
The other three datasets – TMBImmnuo, BreastMSK, LGGGBM – were all sourced from the cBioPortal, https://www.cbioportal.org/ , for cancer genomics.
• GBSG (Rotterdam & German Breast Cancer Study Group) requires prediction of survival time for breast cancer patients.
• METABRIC (Molecular Taxonomy of Breast Cancer International Consortium) requires prediction of survival time for breast cancer patients.
Covariates include expressions for four genes as well as clinical data.
• SUPPORT (Study to Understand Prognoses Preferences Outcomes and Risks of Treatment) requires prediction of survival time in seriously ill hospitalised patients.
Covariates include demographic and basic diagnosis information.
• WHAS (Worcester Heart Attack Study) requires prediction of acute myocardial infraction survival.
• TMBImmuno (Tumor Mutational Burden and Immunotherapy) requires prediction of survival time for patients with various cancer types using clinical data.
Covariates include age, sex, and number of mutations.
Retrieved from https://www.cbioportal.org/ study/clinicalData?id=tmb_mskcc_2018 .
• BreakMSK requires prediction of survival time for patients with breast cancer using tumour information.
Covariates include ER, HER, HR, mutation count, TMB.
Retrieved from https://www.cbioportal.org/study/clinicalData?id=breast_msk_2018 .
• LGGGBM. requires prediction of survival time for cancer patient from clinical data.
Covariates include age, sex, purity, mutation count, TMB.
Retrieved from https://www.
cbioportal org /study/clinicalData? id
= lgggbm tcga_pub
C Further Results
CQRNN
Sequential grid
Excl.
censor
DeepQuantReg
LogNorm MLE
N o r m un
i f
o r m
ObservedCensoredTruequantilesEst.quantile0.1Est.quantile0.3Est.quantile0.
Est.quantile0.7Est.quantile0.
N o r m li n ea r
N o r m non - li n .
E xpon e n ti a l
W e i bu ll
L og
N o r m
[Table]
[Table] [Table]
Difference in means, alongside 95% conﬁdence interval
[Table] [Figure]
C.1 Hyperparameter ablations
The requirements for CQRNN in Algorithm 2 include two hyperparameters unique to CQRNN (the rest determine the NN and its optimisation) – the grid of quantiles, grid τ , and the large pseudo y value, y ∗ .
This section investigates the role of these.
It also discusses two modiﬁcations which were tested in initial experiments but which were not deemed essential for good performance, and excluded from the proposed method – mitigating crossing quantiles and interpolating between quantiles.
Grid size.
For evenly spaced grids, as considered in this work, the number of quantiles estimated,
M , fully determines the grid.
One might expect that a larger grid size, with ﬁner increments between quantiles, would lead to a better ﬁt, since censored weights can be estimated more accurately.
Indeed, the convergence rate of Portnoy’s estimator was found to depend on grid size, O (1 / ( MN ))
[Neocleous et al., 2006].
We hypothesised that a larger grid might only deliver beneﬁt when the dataset was sufﬁciently large for these ﬁne-grained quantiles to be distinguished.
Hence, Figure 4 shows an ablation investigating the interaction between grid size and number of datapoints for each of our type 1 synthetic 1D functions.
Experiments were repeated with 100 random seeds, which was required to reduce the error bars sufﬁciently for comparison.
Grid size was varied, M ∈ { 9 , 19 , 39 } , and number of datapoints, N ∈ { 100 , 200 , 400 , 800 , 1600 , 3200 , 6400 , 12800 } .
For all datasets and grid sizes, TQMSE decreases with dataset size.
In general a larger grid size does produce lower TQMSE, and in three datasets the beneﬁt is signiﬁcant (Norm uniform, Exponential,
Weibull), with the advantage indeed more pronounced with a larger dataset size.
In two datasets (Norm non-linear, LogNorm), this beneﬁt is slight, and a larger grid is even seen to be slightly harmful on very small datasets.
One dataset (Norm linear) does not follow this trend, where the widest grid proves slightly harmful across all dataset sizes.
Pseudo y value.
Portnoy [2003] proposed that y ∗ could be set to any large value approximating ∞ , with the R package ‘ quantreg ’ setting y ∗ = 1 e 6 .
Since the sequential grid algorithm learns quantiles sequentially, it can simply halt if it attempts to estimate a quantile for which only censored datapoints remain, and hence is undeﬁned.
The CQRNN algorithm changes this situation in two ways.
Firstly, since the algorithm is no longer sequential and learns all quantiles simultaneously, it does not have the option of halting.
Secondly, since a non-linear function is learnt, it is possible that higher quantiles might be undeﬁned in one region of the input space, whilst being learnable in the rest of the space.
Regressing towards an arbitrarily large y ∗ value for a portion of the input space could adversely impact the quantile estimate elsewhere.
When all quantiles in grid τ are fully deﬁned, the effect is no different to using ∞ .
To accommodate these differences, we recommend setting y ∗ to a more modest value.
We deﬁne 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 .
In real-world problems, a practitioner might use an estimate for this based on their knowledge about the maximum feasible value for the target.
In lieu of that, we set c y ∗ = 1 . 2 for all our experiments (except the below ablation!), which provided reasonable results across datasets.
Table 6 presents an ablation on four of our (multidimensional) type 1 synthetic datasets, using c y ∗ ∈ { 1 . 0 , 1 . 2 , 1 . 5 , 2 . 0 , 10 , 9 , 100 . 0 } .
For dataset Norm light, y ∗ has no impact since the target distribution is fully deﬁned under censoring.
However, other datasets have input regions where the higher quantiles are not deﬁned due to censoring, and hence higher quantiles are impacted by the magnitude of y ∗ .
The optimal value varies by dataset (e.g. 1 . 2 is best for Norm heavy, while 10 . 0 is best for LogNorm heavy).
[Table]
Crossing quantiles.
One issue often discussed in quantile regression is ‘crossing quantiles’.
Higher quantiles should always produce higher predictions, i.e. ˆ y i,τ 1 > ˆ y i,τ 2 ∀ τ 1 > τ 2 .
The crossing quantile issue arises when this condition does not hold.
We anticipated that the ﬂexibility of NNs might exacerbate this issue, and we tested two methods to remedy this. 1) Adding a crossing penalty to the loss [Bondell et al., 2010], L cross = (cid:80) Ni =1 (cid:80) N τ − 1 j =1 max[0 , c − (ˆ y i, grid τ
[ j +1] − ˆ y i, grid τ
[ j ] )] , where c is the smallest acceptable distance between neighbouring quantiles.
2) Modifying the NN architecture to enforce monotonicity between quantiles, by constraining each consecutive quantile prediction to add on to the previous one, after passing through a SoftPlus .
In our experiments, neither method signiﬁcantly impacted performance.
Favouring simplicity, we propose the CQRNN algorithm without these.
It’s possible that other methods might have more effect, e.g. [Zhou et al., 2020, Brando et al., 2022].
We leave further exploration to future work.
Interpolating quantiles.
The CQRNN algorithm sets the estimated censored quantiles, q , by choosing the prediction closest to the censored datapoint, ˆ q j ← arg min τ | ˆ y j,τ − y j | .
In early experiments, we considered an alternative approach that took a linear interpolation between the two nearest quantiles.
In initial experiments this wasn’t found to signiﬁcantly improve performance, so we propose the CQRNN algorithm using the simpler arg min approach.
Partial vs. full optimisation.
Figure 5 explores the effect of optimising the NNs partially, as is proposed in CQRNN in Algorithm 2, compared to fully, as might be done more typically in EM procedures.
The ﬁgure shows that convergence is fastest when using partial maximisation, when the most up-to-date estimates of ˆ q i are used.
This ends up being more efﬁcient than freezing ˆ q i and only updating after a longer period of optimisation.
[Figure]
[Figure]