Surgical Prediction with Interpretable Latent Representation

1 INTRODUCTION . Hospitals have shown significant interest in exploiting machine learning to predict surgery-related outcomes from postoperative complications ( Xue et al. , 2021 ) to surgery duration ( Jiao et al. , 2020 ) . More than 10 % of surgical patients experience major postoperative complications , leading to increased mortality risk and need for a higher level of care ( Tevis & Kennedy , 2013 ) . Early identification of risk factors can be crucial to early intervention and improved outcomes ( Kable et al. , 2002 ) . Periopertive care is also a major contributor to overall hospitalization expenses ( Childers & Maggard-Gibbons , 2018 ) . Predictions of surgery-related outcomes provides clinicians lead time in resource planning to reduce the cost of periopertive care ( Denton et al. , 2007 ) . However , building a predictive model for perioperative data can be challenging due to the high dimensionality and complexity of perioperative data ( Li et al. , 2020 ) . There are hundreds of perioperative variables of heterogenous types and the missing rates vary drastically across variables , leading to significant deterioration and instability in prediction performance ( LE et al. , 2016 ) . Furthermore , the predictive models need to be interpretable from a clinical perspective for their acceptance and adoption in perioperative care . Deep generative models based on variational auto-encoder ( VAE ) represent a promising approach to deal with the high dimensionality and complexity of perioperative data . Such models can extract the hidden nonlinear relationship of input variables while reducing the input dimensionality through latent representations ( Gao et al. , 2016 ; Kingma & Welling , 2013 ; Rezende & Mohamed , 2015 ) . However , traditional VAEs have two major limitations for surgical predictions . First , as VAE training is traditionally unsupervised , predictive models based on VAE-encoded latent variables can suffer performance degradation ( Feng et al. , 2020 ; Locatello et al. , 2019 ) . Furthermore , the latent space from vanilla VAE models are often entangled ( Burgess et al. , 2018 ; Chen et al. , 2018 ) , which makes it difficult for clinicians to interpret the nonlinear encoding and trust the predictions . Recent advances in VAEs aim to overcome the limitations separately . To improve predictive performance , semi-supervised VAEs are developed to integrate prediction loss into the training of latent encoding ( Pálsson et al. , 2019 ; Cheung et al. , 2014 ; Wu et al. , 2019 ) , but such method does not guarantee disentanglement on latent dimensions other than isolating the specific dimension for prediction . On the other hand , several techniques have been proposed to disentangle the latent space ( Chen et al. , 2018 ; Kim & Mnih , 2018 ; Achille & Soatto , 2018 ; Burgess et al. , 2018 ) , but they are not concerned about optimizing downstream predictions . In fact , existing literature suggests that these disentanglement efforts have little help to predictive tasks ( Feng et al. , 2020 ; Locatello et al. , 2019 ) . To meet the demand for high prediction performance and interpretable models in perioperative care , we propose surgical VAE ( sVAE ) , a representation learning framework specifically designed for high-dimensional and complex clinical data . sVAE provides a holistic framework that integrates the respective strengths of regularization-based disentanglement ( Burgess et al. , 2018 ; Chen et al. , 2018 ) and guided encoding ( Zhao et al. , 2019 ; Feng et al. , 2020 ; Zhou & Wei , 2020 ) . It enhances predictive performance by organizing the latent encoding with respect to the predictive tasks explicitly . Furthermore , it groups patients into interpretable phenotypes in the latent space by disentangling the latent representation . We apply sVAE to two real-world perioperative datasets collected from 12,904 and 77,169 surgeries , respectively . To examine its generalizability to a broader setting , we applied sVAE to the open MIMIC-III clinical dataset and observed similar improvements . Our experiments demonstrate that sVAE consistently delivers superior predictive performance for classification and regression tasks while constructing interpretable latent embedding that captures the intrinsic information of original clinical data and achieves better recovery of original data . While this work is motivated by and evaluated in the context of clinical applications , we note that the proposed approach may be generalized for other fields that demand accurate predictions and interpretable representations based on high-dimensional and noisy data . 2 RELATED WORKS . As a representation learning framework , traditional VAE has been applied to different clinical applications such as abnormality detection ( Baur et al. , 2018 ) or clustering tasks ( Zhao et al. , 2019 ) . Recent research advances VAE in separate directions including semi-supervised models integrating prediction tasks and disentanglement of unsupervised models . Integration of prediction tasks . These exists various types of semi-supervised VAEs that integrate prediction task as a regularizer in latent encoding . Such models , including M2 ( Kingma et al. , 2014 ) , pi-VAE ( Zhou & Wei , 2020 ) , i-VAE ( Sorrenson et al. , 2020 ) and others ( Wu et al. , 2019 ; Pálsson et al. , 2019 ) organize the latent distribution with respect to the prediction ( or conditioning ) variable , hence the prediction information is contained in the representation in latent space . Intuitively the integration of prediction label will lead to easier downstream prediction tasks , as latent representations should have been well-separated by the task . However , as discussed in recent publications ( Feng et al. , 2020 ; Locatello et al. , 2019 ) , a well-trained VAE model does not benefit downstream prediction , often referred to as ” Good ELBO , Bad Inference ” . The extended semi-supervised models adds an explicit loss term to address this issue ( Zhao et al. , 2019 ; Feng et al. , 2020 ) but there are still remaining issues . On one hand , if we assume all latent dimensions to be dependent on the prediction label ( ( Zhao et al. , 2019 ) ) , then as Sorrenson et al . ( 2020 ) point out , the disentangling effects only work when the prediction ( conditioning ) label acts as the cause for the rest latent variables , which obviously does not hold in clinical applications . On the other hand , if we assume the predicted label is independent from the rest of latent dimensions ( Hejna et al . ; Feng et al. , 2020 ) , this assumption is usually violated in model training without a designated constraint . Hejna et al . propose a factor prediction method that emphasizes the independence of prediction dimension from other latent dimensions . However , the encoded position on the prediction dimension is not guided to convey any specific meanings , and the rest of latent dimensions are not disentangled from each other . In contrast , our proposed sVAE maximizes the disentanglement between all latent dimensions and organizes all samples on a prediction dimension based on the predicted label . By doing so , it naturally distills the predictive information to the single prediction dimension and the encoded coordinate can be directly interpretated as predicted values . Unsupervised disentanglement . In β-VAE ( Burgess et al. , 2018 ) , the disentanglement is achieved by directly penalizing KL divergence . In subsequent studies , decomposing this KL term reveals that the total correlation ( TC ) term plays the essential role of disentanglement ( Kim & Mnih , 2018 ; Chen et al. , 2018 ; Achille & Soatto , 2018 ) . As computing TC is intractable , several estimation metrics were proposed ( Kim & Mnih , 2018 ; Chen et al. , 2018 ) , but they underestimate the true TC and lead to unbounded , inaccurate estimation ( Locatello et al. , 2019 ) . In sVAE , we address this issue by evaluating the correlation in the sampled representation and minimizing a bounded metric . Compared to these existing works , sVAE provides a holistic framework combining features of both disentanglement and predictions , enabling the possibility of phenotyping and explicit prediction , thereby meeting the demand for predictive performance and interpretability of surgical models . With more stable estimation of correlation , we hope to achieve both better disentanglement and meanwhile more accurate representation of samples . 3 BACKGROUND . Notations : We denote X ∈ Rn×p as the numerical embedding of the perioperative EHR data , where clinical texts , string categories , and time series have been represented by numerical vectors as illustrated in Appendix Fig . A1 . y ∈ R represents the prediction tasks , either discrete ( classification prediction ) or continuous ( regression prediction ) , and is only available after the surgery . A vanilla VAE consists of an encoder and a decoder , linked by a loss function . The encoder is usually represented by a neural network with parameters ϕ , which outputs qϕ ( z|x ) , a posterior Gaussian probability density of latent variable z ∈ Rd . The decoder is another neural network , denoted by pθ ( x|z ) . Its input is the latent representation z and outputs the probability distribution of the data . As VAE attempts to both reconstruct data and enforce smooth latent distribution , the objective function , usually referred to as Evidence Lower BOund ( ELBO ) , is designed as follows : l ( ϕ , θ ) = N∑ i=1 Ez∼qϕ ( z|x ( i ) ] ) [ log pθ ( x ( i ) |z ) ] −KL ( qϕ ( z|x ) ||p ( z ) ) ( 1 ) The regularizer measures the Kullback-Leibler divergence between the encoded data ’ s distribution and the assumed distribution of z . This regularizer term keeps the latent space sufficiently diverse and meanwhile organize all latent representations conforming to the prior belief . 3.1 INTEGRATION OF TASK VARIABLES . Recent development developed various variants of VAE to incorporate prediction label into the model training . In pi-VAE and similar models ( Zhou & Wei , 2020 ; Khemakhem et al. , 2020 ; Sorrenson et al. , 2020 ) , prediction label is used for conditional distribution to obtain improved disentanglement of latent dimensions . The loss function is defined as : l ( ϕ , θ ) = N∑ i=1 [ −Ez∼qϕ ( z|x ( i ) , ui ) [ log pθ ( x ( i ) |z , ui ) ] +KL ( qϕ ( z|x , ui ) ||p ( z|ui ) ) ] ( 2 ) where u is another observed variable ( or prediction label y in our case ) . A different approach of integrating y is to let y take one of the latent dimensions ( Zhao et al. , 2019 ; Hejna et al . ; Feng et al. , 2020 ) . Depending on the assumptions , the expression of loss function can be different . We refer such approach as smooth-ELBO ( Feng et al. , 2020 ) , defined by : l ( ϕ , θ ) = N∑ i=1 [ −Eqϕ ( z|x ( i ) ) [ log pθ ( x ( i ) |z ) ] +KL ( f ( y ( i ) ) ||qϕ ( y|x ( i ) ) ) ] +KL ( qϕ ( z , y|x ) ||p ( z , y ) ) ( 3 ) where f ( · ) is a smoothing function when y is discrete . 3.2 DISENTANGLEMENT . VAE models are usually highly entangled , making latent space hard to interpret ( Chen et al. , 2018 ; Burgess et al. , 2018 ; Kim & Mnih , 2018 ) . A simple modification by optimizing a heavily penalized KL divergence achieves a high degree of disentanglement in image datasets ( Burgess et al. , 2018 ) . The extra scalar parameter on the KL divergence is β hence this variant is referred as β-VAE . It is not made explicit why penalizing KL ( qϕ ( z|x ( i ) ) ||p ( z ) ) can lead to disentangled transformations , until the ablation study on the decomposition of the KL divergence showed that the total correlation ( TC ) term : TC : = KL ( qϕ ( z ) || ∏d j=1 qϕ ( zj ) ) is the key factor . By emphasizing on the TC term with parameter β in the loss function , we get β-TCVAE and FactorVAE ( Chen et al. , 2018 ; Kim & Mnih , 2018 ) : l ( ϕ , θ , β ) = N∑ i=1 −Eqϕ ( z|x ( i ) ) [ log pθ ( x ( i ) |z ) ] +KL ( qϕ ( z|x ) ||p ( z ) ) + β ∗ TC ( z ) ( 4 ) where zj is the latent representation on dimension j . As TC ( z ) is intractable , several estimation metrics were proposed ( e.g. , density ratio ( Kim & Mnih , 2018 ) , batch sampling ( Chen et al. , 2018 ) ) .

Summary: The paper proposed a modification of the VAE that was designed for handling ICU data for perioperative care. The new method specifically handles the entanglement of the latent variables in the traditional VAE, so it can provide interpretable representation of the data that also serves well the downstream prediction and regression tasks. The authors conducted experiments on one publica dataset and two new real-world perioperative datasets. In all cases, the proposed methods are superior to the comparator methods.