Towards modelling hazard factors in unstructured data spaces using gradient-based latent interpolation
Abstract
The application of deep learning in survival analysis (SA) allows utilizing unstructured and high-dimensional data types uncommon in traditional survival methods.
This allows to advance methods in ﬁelds such as digital health, predictive maintenance, and churn analysis, but often yields less interpretable and intuitively understandable models due to the black-box character of deep learning-based approaches.
We close this gap by proposing 1) a multi-task variational autoencoder (VAE) with survival objective, yielding survival-oriented embeddings, and 2) a novel method HazardWalk that allows to model hazard factors in the original data space.
HazardWalk transforms the latent distribution of our autoencoder into areas of maximized/minimized hazard and then uses the decoder to project changes to the original domain.
Our procedure is evaluated on a simulated dataset as well as on a dataset of CT imaging data of patients with liver metastases.
Introduction
Survival analysis (SA) is indispensable for estimating factors of increased risk, e g., for an infection of a contagious infection or the onset of a disease.
Survival is also an important clinical endpoint for assessing safety and efﬁcacy of cancer therapies [1].
Despite its importance and practical relevance, modelling survival times remains challenging due to the often complex data generating process, involving censoring, truncation, time-varying features, recurrent events or competing risks.
While there exist frameworks such as [2, 3] that reduce the complexity of modelling survival related tasks by relating SA to a Poisson regression using a special data transformation [4], many recent approaches adapt the semi-parametric Cox proportional hazard (PH) model [5].
Instead of relating directly to the survival probability, the Cox PH models the hazard rate, which is the instantaneous risk of experiencing an event (disease, death, etc.).
Modern non-linear adaptions of the Cox PH model employ neural networks and advanced optimizers [6, 7].
Other approaches strive to model the hazard at every possible point in time using a discretized output time domain [8, 9, 10].
More recently, neural ordinary differential equations for use in SA were applied successfully [11].
Generative models and survival have been investigated in [12], who use an autoencoder in a semi-supervised setting to encode survival information and classify between long- and short-term survivors.
Moreover, [13] regularize a denoising autoencoder with a survival loss to perform cardiac motion analysis.
SA is predominantly applied for tabular data.
A few exceptions exist (see, e.g., [14]), but a majority of techniques is not designed for unstructured or high-dimensional data.
Modern deep learning approaches address this problem, but in contrast to existing SA techniques, do not have a straightforward model interpretation – a crucial aspect in digital health and medical decision making.
This interpretation and understanding is equally important for generative models, yet more intricate and especially challenging in the light of SA.
We close this gap by proposing a two-step procedure combining ideas of deep generative models and counterfactuals [15].
In a ﬁrst step, we estimate an inherently interpretable latent representation of the unstructured data source using a multi-task variational autoencoder (VAE; [16]) with regularizing survival objective.
In the second step we help the practitioner to understand the modelled relationship of this latent space and the survival outcome using our novel HazardWalk method.
While our autoencoder architecture allows modelling and isolating survival-relevant factors, the HazardWalk can be used to traverse in the latent space guided by the survival-objective and utilizes the generative decoder to project back its reasoning to the original data space (e.g., by an respective visualization in the image).
Methods
Autoencoder with survival downstream objective Our proposed architecture consists of a probabilistic encoder enc φ and a decoder dec θ with parameters φ and θ , respectively.
For a sample x ∈ X from the sample space X , the approximate inference distribution q φ ( z | x ) of the encoder is given by
N ( µ φ , diag ( σ 2 φ )) where µ φ , σ φ = enc φ ( x ) .
By sampling z ∼ q φ ( z | x ) , a reconstructed version ˆ x of the data sample x can be produced by dec θ .
In order to learn this reconstruction, we optimize the evidence lower bound (ELBO)
[Equation]
where log p θ ( x | z ) accounts for the reconstruction and β is an additional parameter to control the impact of the KL-divergence KL between q φ ( z | x ) and the Gaussian prior p ( z ) [17].
While the plain ( β -)VAE works well for reconstruction, it does not learn concepts to emphasize survival-relevant information in x .
We therefore propose to extend the ( β )-VAE to a multi-task model with survival downstream objective by adding a network head cox ψ with parameters ψ to the latent bottleneck.
This leads to a latent space guided by survival information.
The task of cox ψ is to predict the log-hazard rate or risk score r ∈ R with r = cox ψ ( z ) for z ∼ q φ ( z | x ) .
In SA, the hazard h ( t | x ) describes the instantaneous risk of experiencing the event of interest exactly at a speciﬁc time point t .
Following the original Cox-PH model, the hazard can be decomposed as h ( t | x ) = h 0 ( t ) exp( r ) , where h 0 ( t ) is a feature-independent baseline hazard and exp( r ) serves as a feature-dependent scaling factor based on the risk score r .
Let X be a dataset of n observations x i with corresponding event times t i for i = 1 , . . . , n .
For a predicted risk score r i =
cox
ψ ( z i ) based on z i ∼ enc φ ( x i ) , the corresponding objective function is given by the negative partial Cox log-likelihood (see, e.g., [6]):
[Equation]
where δ i = 0 if the i th observation is censored, else δ i = 1 .
By regularizing the autoencoder with the survival downstream task, latent features are forced to learn a representation that is inherently interpretable from a survival point of view.
The joint objective of our multi-task approach is given by
[Equation]
with τ ∈ [0 , 1] controlling the inﬂuence of each loss term in the convex combination of both.
In our experiments, a value of τ = 0 .
5 turns out to give both good reconstructions and survival-oriented latent spaces.
Hazard guided walk In order to visualize the impact of survival-relevant dimensions in the latent space, we propose the gradient-based hazard latent walk ( HazardWalk ).
HazardWalk is a new method to traverse the latent space based on the estimated hazards by cox ψ and utilizes dec θ as generative model to project these changes to the original data space, making important factors on the overall survival directly observable.
The latent variational distribution of the VAE thereby ensures smooth transitions in X .
First, the latent feature distribution is obtained by forwarding a sample x through enc φ , yielding parameters ( µ φ , σ φ ) = ξ .
The goal is now to obtain the expected gradient of the hazard rate w.r.t.
the whole latent distribution, which points in the direction of maximal hazard rate increase.
We approximate this using Monte Carlo sampling with
[Equation]
based on samples z
b ∼ q φ .
By updating the parameters of q φ ( z | x ) in the direction of (4) and repeating this procedure for a pre-speciﬁed amount of iterations, we shift and scale the latent distribution into areas of maximized hazard.
We can then use the decoder dec θ to produce new samples (cid:101) x from this transformed latent distribution.
In this case, the disentanglement of latent features induced by the factorized Gaussian prior leads to samples, where survival-irrelevant characteristics are left relatively unchanged and survival-relevant characteristics are changed notably.
In other words, the
HazardWalk allows practitioners to understand the model’s learned relationships directly in the original sample space by, e.g., plotting images based on the most survival-relevant changes in the latent space.
This procedure can also be reversed, using the negative expected gradient for the parameter updates to generate samples with minimal risk.
Experiments
We evaluate the HazardWalk procedure on two datasets.
First, on a simulated dataset based on
MNIST, for which survival times are generated as described in [18].
The result is an MNIST dataset with the original images, where high hazard rates correspond to a high digit and vice versa.
Second, we demonstrate the proposed approach on a dataset of 492 CT scans of patients with liver metastases, where the survival time after imaging is known.
Our enc φ and dec θ each consist of four residual blocks [19] with downsampling or upsampling operations respectively.
cox ψ was chosen to be a fully-connected network with no hidden layer and no bias, which equals a linear predictor.
We use Adam [20] as optimizer with a learning rate of 1 efor θ, φ and 1 e- 5 for ψ .
HazardWalk on MNIST Based on a latent space of 4 dimensions, Figure 1 shows the differences in the embedding of a VAE with and without the regularizing Cox objective on the test dataset.
As expected, the vanilla VAE (left plot) does not preserve the natural order of MNIST numbers, but structures the latent space with the intention of maximum visual reconstruction, hence the close neighborhood of digits like 0 and 9 .
In contrast, when adding the survival information (right plot), the behavior changes.
The main source of variation is now based on the actual hazard, i.e., the natural digit order.
Digits such as 0 and 1 are now placed next to each other instead of being maximally separated as for the vanilla VAE due to their visual dissimilarity.
In a second step we investigate the effect of the HazardWalk on these survival-oriented embeddings.
We apply the HazardWalk with 1500 iterations, 128 samples and visualize the progress of the transformation.
As depicted in Figure 2, the gradients of cox ψ lead to a transition into different digit [Figure] classes with higher (or lower) hazard.
It is noteworthy that each digit shifts into a class that allows to keep some of the original core features (e.g., a 2 turns into a 3 ﬁrst, then into an 8 instead of directly transforming into a 9 with the maximum associated hazard rate).
HazardWalk on computed tomography imaging data CT imaging is an important step in the diagnostic workup of cancer patients and contains prognostic information.
Here, we used axial, contrast-enhanced portal-venous CT images of 492 patients with hepatic metastases of colorectal cancer.
Images were resampled to 1 mm 3 isotropic spatial resolution and liver as well as hepatic tumors were segmented using a pretrained nnU-Net [21]. For further computations, the volume was centered on the liver segmentation, cropped to 240 3 voxels and downscaled to a resolution of 64voxels. On this volume, the Cox-regularized VAE was trained with a latent space of 8.
The input to the network consists of the actual masked CT data and the segmentation of the tumor, represented as two channels.
A major challenge for model training is the small sample size in combination with a relatively high dimensional input space.
This results in gaps in the otherwise smooth latent space.
By increasing [Figure] the β -parameter and therefore the inﬂuence of the KL -term, this can be mitigated at the cost of lower-quality reconstructions.
Figure 3 depicts multiple examples, showing different liver cancer cases where the hazard ratio is tweaked by HazardWalk.
Our method conﬁrms the meaningfulness of the learned latent space of the autoencoder, increasing number and sizes of tumor patches when increasing hazard.
The shape of the liver changes only slightly, which indicates that overall survival may be unrelated to this characteristic.
In the given example, the sample in the ﬁrst row of Figure 3a has a time-to-event of 1479 days and an estimated hazard rate of 0.601.
This constitutes a hazard that is signiﬁcantly lower than the baseline.
After the transformation, the newly obtained (cid:101) x is estimated with a hazard rate of 3.424.
On the contrary, reducing the hazard leads to exactly the opposite – the tumor shrinks or vanishes.
This can be seen, e.g., for the ﬁrst sample in Figure 3b, with an original time-to-event of 61 days and a hazard rate of 3.165 that transformed to 0.785.
Conclusion and Outlook
In this work we proposed a novel two-step procedure to learn and interpret the relationship between unstructured data sources such as images and survival times.
While our generative model allows to emphasize hazard factors in unstructured data spaces, the proposed method HazardWalk yields a straightforward interpretation of the model using a gradient-based latent interpolation and reconstruction over the generative decoder.
One limiting factor in training generative models for SA are the typically small-scale datasets.
This results in gaps in the latent space and/or insufﬁcient representations.
Adjusting the KL -term turned out to be beneﬁcial in our experiments to focus on the quality of the representation, resulting in more reliable and stable HazardWalk interpolations.
While there are larger CT datasets available (e.g., [22]), these are mostly for segmentation tasks and not associated with a survival outcome.
An extension of our approach to a semi-supervised setting could thus allow the integration of more information to improve the encoding and generation.
Alternatively, adversarial autoencoders as in [23, 24] could help to improve the quality of generated images.
Given additional tabular data, a possible direction for the VAE’s downstream task could also be a semi-structured model [25].
Acknowledgments
This work has been partially supported by the German Federal Ministry of Education and Research (BMBF) under Grant No. 01IS18036A.
We thank the anonymous reviewers for their constructive comments, which helped us to improve the manuscript