S URGICAL P REDICTION WITH I NTERPRETABLE L ATENT R EPRESENTATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Given the risks and cost of surgeries, there has been significant interest in exploiting predictive models to improve perioperative care.
However, due to the high dimensionality and noisiness of perioperative data, it is challenging to develop accurate, robust and interpretable encoding for surgical applications.
We propose surgical VAE (sVAE) , a representation learning framework for perioperative data based on variational autoencoder (VAE).
sVAE provides a holistic approach combining two salient features tailored for surgical applications.
To overcome performance limitations of traditional VAE, it is prediction-guided with explicit expression of predicted outcome in the latent representation.
Furthermore, it disentangles the latent space so that it can be interpreted in a clinically meaningful fashion.
We apply sVAE to two real-world perioperative datasets and the open MIMIC-III dataset to evaluate its efficacy and performance in predicting diverse outcomes including surgery duration, postoperative complication, ICU duration, and mortality.
Our results show that the latent representation provided by sVAE leads to superior performance in classification, regression and multi-task predictions.
We further demonstrate the interpretability of the disentangled representation and its capability to capture intrinsic characteristics of surgical patients.
While this work is motivated by and evaluated in the context of clinical applications, the proposed approach may be generalized for other fields using high-dimensional and noisy data and valuing interpretable representations.
I NTRODUCTION
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 (P´alsson 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, 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.
R ELATED W ORKS
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; P ´ alsson 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.
B ACKGROUND
Notations: We denote X ∈ R n × 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 ∈ R d .
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:
[Equation]
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. I NTEGRATION 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:
[Equation]
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 [Equation] the expression of loss function can be different.
We refer such approach as smooth-ELBO (Feng et al., 2020), defined by:
[Equation]
where f ( · ) is a smoothing function when y is discrete.
3. D ISENTANGLEMENT
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 ) || (cid:81) dj
=
1 q ϕ ( z j )) 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):
[Equation]
where z j 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)).
S VAE
The aforementioned works have significant implication for surgical predictions.
Conceptually, disentanglement helps reorganize the information in X into clinical characteristics/phenotypes that we may be able to discover; and the integration of task variables enable VAE to stratify surgical cases by predicted labels/values that we are interested in.
Hence with such goals, we want our sVAE to 1) express explicit predictions on a latent dimension (serving as a predictor); 2) minimize the correlation between latent dimensions so that the latent dimension demonstrates interpretable phenotype.
If we have an accurate explicit prediction as in 1), we would subsequently hope that it also has an accurate representation for other phenotypes, hence reconstructs the samples robustly. Distillation of Prediction.
Now consider a single case x ∈ X : as the prediction task y is dependent on x (otherwise the problem is not meaningful), it is also dependent on its latent encoding z .
This fact can be represented as: p ( y, x , z )
=
p ( y | z , x ) p ( z | x ) p ( x ) .
To distill the prediction information y from z , let g ( · ) be a disentangling transformation function such that g ( z ) =
[ y, z ′ ] , where z ′ is a vector of the remaining information in z that is irrelevant to y .
As g ( · ) is a disentangling function, all dimensions in z ′ are uncorrelated with each other too.
This is similar to an independent component decomposition except that we explicitly isolate y from other dimensions.
Let v =
[ y, z ′ ] , then p ( v ) = p ( y ) p ( z ′ ) .
Since v is a transformation of z , it is equivalently another latent representation of x .
Assuming g ( · ) is bijective, then p ( y, x , z ) can be re-written as p ( y, x , z )
=
p ( y, x , v ) =
p ( v | y, x ) p ( y | x ) p ( x ) = p ( z ′ | x ) p ( y | x ) p ( x ) .
We introduce g ( z ) for two purposes: it both forces the disentanglement in all latent dimensions for potential interpretability, and distills the predictive information from z and explicitly expresses it in a designated dimension.
We use a neural network to implement g ( z ) , then its generative version g ( q ϕ ( z | x ( i ) )) is equivalently another encoder.
Let g ( q ϕ ( z | x ( i ) ))
=
q ϕ ′ ( v | x ( i ) ) , now we leverage the idea in Eq.3 for the learning of p ( z ′ | x ) and p ( y | x ) , and the TC term in Eq.5 for our assumption of disentanglement to hold true:
[Equation]
where α and β are the scalar coefficients for prediction loss and TC loss.
Note that the smoothing function f ( y ( i ) ) is unnecessary, especially when y is a regressional task.
In our experiments, we replace it by the posterior Gaussian likelihood log q ϕ ′ ( y ∗ ( i ) | x ( i ) ) where y ∗ ( i ) is the ground truth label, and set hyperparameters α = 1 and β = 10 .
Compared to other prediction-guided models, the introduction of TC term not only guarantees the conformation to our independence assumption, but also eliminates the use of an additional predictor as prediction is explicitly expressed on a single dimension.
Estimation of Correlation .
Recent studies show that the estimation of TC is usually misleading.
As TC is measured by a KL-divergence from batch data B , TC = KL ( q ϕ ( z ) || (cid:81) dj
=
1 q ϕ ( z j )) ≈
[Equation]
ϕ
[Equation] [Equation] j ) . When d is small, such approximation is appropriate.
When d is large, there are several issues.
First, when sampling z ∼ q ϕ ( z ) , due to the curse of dimensionality, the input sparsity leads the estimation q ϕ ( B ) ( z ) effectively 0, hence the estimation of such
KL-divergence is insignificantly different from 0 and will not disentangle the latent space anymore (Kim & Mnih, 2018). In real applications, we may take the log term of the estimated KL divergence, and this inappropriate estimation would impose a negatively large, unbounded fluctuation in the loss minimization. Second, the estimation of TC using such method with mean representation is not consistent with the sampled representation (Locatello et al., 2019).
Although several modified estimation methods were proposed (Kim & Mnih, 2018; Chen et al., 2018), with a higher dimensionality, these are no longer accurate and further lead to unintended shutdown of latent dimensions (Cheng et al., 2020). Besides, with more training iterations, the posterior variance of q ϕ ′ ( v | x ) is usually smaller, amplifying such negative effects (see experiment results in Appendix A6).
To tackle this issue, we propose a simple estimation of disentanglement using singular value decomposition (SVD).
The intuition behind our method is as follows: if the latent dimensions are disentangled, a linear transformation to a new coordinate system will not increase the variance on any new coordinates.
In fact, it has been observed in medical imaging data that PCA/SVD-based metrics are equivalently effective as TC-based metrics in quantifying similarities (Guyader et al., 2016), which inspires the substitution of the inaccurate TC estimation by SVD.
The evaluation of our proposed metric consists of three steps:
Draw Monte Carlo sampling from the latent distribution of current batch data (batchsize := B > d ) .
Concatenate into a B × d matrix A =
[ v 1 , v 2 , ..., v B ] T .
As each latent dimension may have different mean/variance, each dimension should be standardized.
Perform SVD to get the largest singular value σ 1 , where σ 1 ≥ σ 2 ≥ ... ≥ σ d ≥ 0 .
Calculate TC ( v ) =
σ 21 − 1 d − 1 (similar to the variance explained by the first principal component).
A main advantage of this correlation estimation method is the lack of hyperparameters or expensive computation, which should provide more stable training in an efficient way.
Besides, our proposed evaluation strategy has 3 key properties:
TC ( v ) is bounded by [0 , 1] ; 2.
TC ( v ) is minimized iff the all latent dimensions are uncorrelated (orthogonal); 3. TC ( v ) = TC ( v (1) , v (2) , ..., v ( B ) ) estimated from a single batch data is the lower bound of the TC of the whole latent space, calculated by TC ( v (1) , v (2) , ..., v ( N ) ) .
Meanwhile, the minimization across all batches minimizes the upper bound of the whole latent space
The detailed proofs are discussed in Appendix A1.
E XPERIMENTS
We evaluate sVAE on the open MIMIC-III dataset and two real-world perioperative datasets in terms of its perdictive performance and its efficacy in disentangling latent space and representing the true information.
We further show how disentanglement enhances the interpretability of latent encoding from a clinical perspective.
5. D ATASETS
Delirium Data : The first dataset was acquired from the electronic anesthesia records of all adult patients undergoing surgery at a large academic medical center from June 1, 2012 to August 31, 2016.
Input variables include demographics, history of comorbidities, lab tests, medications, and statistical features extracted from time series.
The prediction outcome is binary, indicating whether delirium was developed after surgery.
In total, there are 12,904 patients with 52.65% of them developed postoperative delirium.
This dataset is referred hereafter as Delirium Data .
OR Data : The second dataset was acquired from a large healthcare system spanning across eight hospitals, academic and community centers and surgical specialties from March 1, 2019 to October 31 2019.
The input variables include the scheduled surgery duration, clinical notes describing [Table] surgery type, and categorical strings (surgeon name, anesthesiologist name and location of the operating room).
23 features were extracted as input from 77,169 surgery cases.
The prediction outcome for this dataset is the actual surgery duration.
This dataset is referred hereafter as OR Data .
In addition, we also evaluate sVAE on the open MIMIC-III (Johnson et al., 2016) for two prediction tasks: a regression task to predict the ICU length of stay (days) and a classification task to predict mortality.
Due to space limitation, the supplementary results on the three dataset are reported in the Appendix A3 and A4.
5. B ASELINES
We compared sVAE against state-of-art latent encoding and data transformation methods.
For deep generative models, we have included VAE , pi-VAE that organizes latent space by conditioning on prediction outcome (Zhou & Wei, 2020), smooth-VAE (Feng et al., 2020) that represents a broad class of VAEs that express predicted outcome as a latent dimension using the smooth-ELBO objective function in Eq.3, β -TCVAE which is an unsupervised disentangled VAE.
Common encoding/transformation techniques such as autoencoder ( AE ), principal component analysis ( PCA ), and independent component analysis ( ICA ) are also implemented.
For fair comparison, all neural network-based models are built with same complexity (same number of layers and number of nodes in each hidden layer).
As dimension reduction is one of our objective, data were compressed to a lower-dimensional representation ( d = 10 ) for all methods.
To evaluate the prediction performance using the encoded/transformed data, several linear and nonlinear predictors are built including logistic regression (LR), support vector machine (SVM) with radial basis function kernel, XGBoost, and deep neural network (DNN).
The hyper-parameters and implementation are fixed throughout the study and details are described in Appendix A4.
5. E VALUATION M ETRICS
We use the following metrics to evaluate the downstream prediction performance, latent disentanglement and input reconstruction.
In each experiment, we sample each dataset 5 times randomly and report the mean and standard error of each metric.
Downstream Prediction: For classification predictions, we report the area under the ReceiverOperating Characteristic curves (ROC AUC) and the average prediction (equivalent to the area under the Precision Recall curve).
The regression performance was evaluated by the root mean square error (rMSE) and R 2 .
As both smooth-VAE and the proposed sVAE explicitly express the predicted outcome as a separate latent dimension, the encoded position on this dimension can be used directly for prediction, by viewing the coordinate on that latent dimension as the predicted values.
This is referred as Direct Prediction hereafter.
Latent Disentanglement: Previous studies (Kim & Mnih, 2018; Chen et al., 2018; Burgess et al., 2018) proposed various metrics for quantifying disentanglement, assuming that we know the factors of variation in the underlying true generative model.
This is infeasible in clinical studies.
Instead, we use the recently developed Unsupervised Disentanglement Ranking (UDR)
(Duan et al., 2019) to assess disentanglement effects.
This method relies on the assumption that for a particular dataset and a VAE-based disentangled representation class, disentangled representations are all alike, while every entangled representation is entangled in its own way.
Therefore, it can be useful when comparing the latent encoding of models with the same dimensionality, but the metric cannot be used to compare models with different latent dimensionality.
When using UDR, higher value implies better disentanglement.
Input Reconstruction: Besides prediction performance and disentanglement, it is essential to have a true representation of the original data, such that the disentanglement and interpretation make sense.
To evaluate the representation accuracy, we reconstruct the input by decoding the sampled representations from the latent space to the original input space, and the reconstruction error is evaluated by the negative log-liklihood (NLL) of getting the original input from the reconstructed input. A lower NLL implies better accuracy in the reconstruction.
[Figure]
a. Delirium Data
[Figure]
b. OR Data
- cation prediction (Fig 1.a), cases with positive postoperative delirium are shown in red and cases with negative postoperative delirium are shown in blue. In regression (Fig 1.b), cases are divided into lower quartile (cyan inter
(green and upper quartile (blue predicted/actual surgery duration [Figure] [Figure] [Figure]
L a t e n t D i m e n s i o n
Q1Q
a: Encoding of albumin levels
[Figure]
Separating intraabdominal surgeries from other types (OR Data)
[Figure]
abdomencolonlaparoscopiccolonoscopyesophagogastroduodenoscopyjoint
[Figure]
c
[Figure]
surgeries (OR
[Figure]
5. Q UANTITATIVE C OMPARISONS
sVAE has better prediction performance.
When fixing latent dimensionality to d = 10 , we observe that sVAE produces consistently better ROC AUC and Average Precision in Delirium data [Table] [Table] (Table 1) and lower rMSE (higher R 2 ) in OR data (Table 2).
The prediction performance is even better than predicting on the full-dimensional raw data ( d = 562 in delirium and d = 23 in OR data).
Compared to PCA or ICA, encoding/disentangling the raw data ( β -TCVAE) does not benefit the downstream tasks, consistent with previous results by Locatello et al.
Compared to unsupervised AE/VAEs, pi-VAE has slightly better prediction performance, but much lower than models that introduce prediction loss into the loss function (smooth-VAE and sVAE).
sVAE does not require a predictor.
sVAE distills the predictive information into a single designated dimension and force the other dimensions to be disentangled from this dimension.
Consequently, adding a predictor has little improvement over the direct prediction provided by sVAE.
This designated prediction dimension is visualized in Fig.
For Delirium Data (Fig. 1.a), we force data stratification on latent dimension 0 with respect to the binary outcome of postoperative delirium.
As the latent space is generally smooth, the stratification can be viewed analogous to the predicted risk of getting delirium.
For OR Data (Fig. 1.b), data are stratified with respect to the predicted surgery duration.
As we minimize the error between prediction and actual duration, the observed distribution on dimension 0 is close to the actual duration.
Consistent results are obtained on other datasets (see Appendix A5).
sVAE has consistently better disentanglement.
To evaluate the disentanglement effects of the proposed correlation estimation metric, we replaced the TC ( v ) by the TC term in β -TCVAE and the Minibatch-Weighted Sampling proposed by Chen et al.
(2018), which is denoted as sVAE TC .
We compared sVAE
TC with sVAE with various dimensionalities.
We use Delirium Data as its original input dimensionality is high ( d = 562 ), giving us flexibility in adjusting the latent dimensions.
We vary the latent dimensionality from 5 to 30.
Besides, a full-dimensional latent space ( d = 562 ) is also created for reference.
At each dimension, we calculate UDR A2A (Duan et al., 2019) for pairwise comparison, and observed that sVAE has consistently better disentanglement than sVAE
TC at every dimensionality (see Table 2).
As UDR A2A is computationally expensive, the calculation fails for the full-dimension model.
To compare with existing baseline VAE models, we have calculated UDR A2A in three datasets and observed that sVAE is consistently better than other VAE models (Table 3).
Table 3 also demonstrates the inaccurate estimation of TC when d = 10 , hence β -TCVAE does not work well in disentanglement when compared to pi-VAE.
Meanwhile, it is noteworthy that the UDR metrics in sVAE and smooth-VAE should not be compared directly with other baseline models, as they have a explicit latent dimension that conveys prediction information, equivalent to a supervised disentanglement for this dimension.
sVAE has consistent prediction and reconstruction performance with dimension reduction.
Due to the bounded values of TC ( v ) , sVAE is able to minimize other loss terms, achieving bet- ter distillation of predictive information and better input reconstruction.
When varying the latent dimensionality, the prediction and reconstruction performance is stable.
On the other hand, the wrong estimation of TC term in sVAE TC introduces instability in reconstruction, leading negative log-liklihood (NLL) to 0.4916 when d = 25 .
In extreme cases ( d = 562 ), the unbounded TC loss dominates the model training and NLL deteriorates to 1.314.
5. I NTERPRETABILITY THROUGH D ISENTANGLEMENT
We now demonstrate through examples that disentanglement enhances the interpretability of the latent representation, which may help clinicians gain trust of the model.
Specifically, we show that the model exhibits phenotyping effects (Basile & Ritchie, 2018), capturing the intrinsic characteristics of patients from the noisy input data, showing patterns in agreement with the existing clinical evidence.
Here we give two examples from Delirium Data and OR Data, respectively.
Examples of
MIMIC-III are in Appendix A6.
In Delirium Data, the input variable albumin level is noteworthy, as clinical evidence showed that it is associated with postoperative delirium (Zhang et al., 2018).
If sVAE captures the intrinsic association between albumin and delirium, then dimension 0 (representing the delirium prediction) should contribute to the distribution of albumin levels in the latent space.
As shown in Fig.
2a, the distribution of albumin values are controlled by both dimension 0 and another dimension, in agreement with clinical expectation.
In OR Data, the most complex variable is planned surgery description consisting of free texts with varying lengths, formats and unsystematic expressions.
It usually requires human efforts to read and interpret the notes before billing and staffing.
Among the surgical description texts, the most frequent surgery types are various types of intra-abdominal surgeries ( abdomen, colon, esophagogastroduodenoscopy, laparoscopic, and colonoscopy ), occurring in from 16.75% to 4.16% of test cases.
Besides, eye-related surgery constitutes 16.75% and bone/joint related keywords ( knee, leg, spine, anthroplasty ) constitutes 9.54% in total.
The various types of intra-abdominal surgeries are all performed on soft tissues in abdomen regions.
They generally have different indications, recovery profiles, and complications from orthopedic or ophthalmic surgeries.
They usually also last longer than orthopedic or ophthalmic surgeries.
In addition, the inclusion of surgeon name as an input variable provides another reason for there to be a latent class dividing surgeries based on the part of the body, as most surgeons specialize in operating on a particular organ.
In the latent space, we observe that latent dimension 5 represents the characteristics of intra-abdomenal surgeries, and separates them from eye-related or joint/bone surgeries (Fig. 2b).
We further split abdomen surgeries with respect to specific keywords, with each circle representing the cluster of a keyword and radius representing the average distance between cases and the cluster center.
As shown in Fig 2.c keywords representing similar surgery types/regions related to abdomen are clustered to the right, with little overlap to eye or joint/bone related surgeries.
This implies that latent dimension 5 learns the intrinsic information underlying the different textual expressions, when different keywords are used to describe intra-abdomen regions.
C ONCLUSION
In this paper, we present sVAE, a prediction-based, interpretable deep latent model.
sVAE provides a holistic framework that integrates regularization-based disentanglement and prediction-guided encoding.
We apply sVAE to two real-world perioperative datasets and evaluate its efficacy and performance in the context of surgical applications.
Our results demonstrate sVAE’s latent prediction dimension can be used directly for prediction with superior performance over predictors based on state-of-the-art VAE models.
Furthermore, the disentanglement of latent spaces demonstrate the phenotyping effects by capturing the intrinsic characteristics of clinical data, and create clinically meaningful organizations from the noisy inputs.
While sVAE was motivated by clinical applications, the proposed approach makes general advancement toward future adoption of VAE models in fields where predictive performance and interpretability are essential in the presence of high-dimensional and noisy data.
E THICS S TATEMENT AND R EPRODUCIBILITY
There are three datasets involved in this application study.
Delirium Data and OR Data were acquired with appropriate IRB approval of the study with waivers of consent.
These clinical datasets were not deidentified or approved for publication.
For reproducibility, the same experiments were repeated on an open clinical dataset MIMIC-III (Johnson et al., 2016) and the corresponding results were attached in Appendix A4-5.
Detailed description of model hyper-parameters and Python code repository can be found in Appendix A3.
A A PPENDIX A. E STIMATION OF C ORRELATION
Generate a latent sample from v ( i ) ∼ q ϕ ′ ( v | x ( i ) ) , where i = 1 , ..., B are the indices of samples of current batch data.
Standardize and concatenate into a B × d matrix A
= [ v 1 , v 2 , ..., v B ] T .
As each latent dimension may have different mean/variance, each dimension should be standardized, yielding A .
Perform singular value decomposition, we get A = U Σ V T = (cid:80) di
=
σ i u i v Ti , where σ 1 ≥ σ 2 ≥ ... ≥ σ d ≥ 0 .
Then TC ( v ) has 3 key properties:
TC ( v ) =
σ 21 − 1 d − 1 is bounded by [0 , 1] ; 2.
TC ( v ) is minimized iff the all latent dimensions are uncorrelated (orthogonal); 3. TC ( v ) =
TC ( v (1) , v (2) , ..., v ( B ) ) is the lower bound of the TC of the whole latent space, calculated by TC ( v (1) , v (2) , ..., v ( N ) ) .
Meanwhile, the minimization across all batches minimizes the upper bound of the whole latent space.
A.1. P ROPERTY
To see 1, note that A is standardized therefore by definition we get || A || 2 F = (cid:80) di =1 (cid:80) nj
{+=
+}
A 2 i,j = d .
Besides, it is easy to show that |
[Equation]
F = d (cid:88) i =
σ 2 i (6) which implies that for the largest singular value we have 1 ≤ σ i ≤ √ d .
A.1. P ROPERTY
To see 2: • when TC ( v ) is minimized, σ 1 = 1 .
As σ 1 is the largest singular values and by Eq. 9, we get σ 1 = σ 2 = ...
=
σ d = 1 .
Therefore, A = U Σ V T = UV T .
Since A is now a product of orthogonal matrices hence itself is also orthogonal.
• In the reverse direction, if all latent dimensions are uncorrelated, we get AA T = 1 , so the singular values of A are all 1 and TC ( v ) = 0 .
For the property 2, intuitively TC ( v ) measures the maximum variance that can be explained by a principal component.
When the largest principle component can only explain minimum variance, we know that the latent space is decorrelated (disentangled).
A.1. P ROPERTY
To prove 3, we first show that TC ( v ) is the lower bound. Note that the singular values of A are square roots of the eigenvalues of A T A . When estimating the whole latent space ( B = N ), this is equivalent to
[Equation]
where X are the latent variables of data not in the current batch ( X contains N − B rows and d columns).
So the question is equivalent to estimating the eigenvalues of the sum of symmetric
(in this case also positive semidefinite) matrices using only A T A .
From Courant-Fischer min-max theorem, we know that the eigenvalues of A T A are:
[Equation]
where R A T A ( x ) =
( A T Ax , x ) ( x , x ) .
Since R A T
A + X T X ( x ) ≥ max { R A T A ( x ) ,
R X T X ( x ) } , we know that λ k
( A T A ) ≤ λ k
( A T A +
X T X ) , hence TC ( v ) forms the lower bound of the true TC ( v ) when estimated using whole data.
Next, we show that after each epoch, we have minimized the upper bound of the actual eigenvalues of the true data.
Denote the first batch of data by A 1 , and then the whole data can be expressed as
[Equation]
, we know that D T D =
A T 1 A 1 + X T 1 X 1 .
Then
[Equation]
hence
[Equation]
After a complete epoch looping through all training data, we have actually minimized the upperbound the the true estimation.
When the training converges, we would know that max { TC i ( v ) } T 1 ≤
TC ( D ) ≤ sum { TC i ( v ) } T 1 .
A. G RAPHICAL DIAGRAM OF S VAE
As raw EHR data consists of non-numerical or non-static variables, a standardized data-processing was conducted for each dataset to get the static, numerical representation of X as illustrated in Fig.
A1.
A. H YPER - PARAMETERS AND IMPLEMENTATION DETAILS
All deep latent models, including AE, VAE, β -TCVAE,smooth-VAE, and sVAE are implemented using Python Pytorch (Paszke et al., 2019), with the same configuration for all datasets and models.
Essential parameters include: activation function =
exponential linear unit (ELU), optimizer =
Adam (learning rate= 1 e − 6 ), batchsize = 256, number of hidden layers = 4 (same for both encoder and decoder), and number of hidden units (10*d, h, 4*h, h) for decoder and (h, 8*h, 4*h, h) for encoder, where h is a hyper-parameter.
Currently h
= 2000.
The DNN predictor is also using Pytorch with the same set of parameters as every encoder.
Other predictors (LR, XGBoost and SVM) and transformation methods (PCA and ICA) are implemented by Sklearn package Pedregosa et al.
For LR, we used l1 penalty and linear solve with regularzier C=0.2.
For Xgboost, the number of estimators for XGBoost is 10 (same as the latent dimension) with loss as ”ls” in regression and ”binary:logistic” in classification.
LogisticRegression(penalty=’l1’,solver=’liblinear’,C=0.2).
For SVM, the ROC AUC and average precision are calculated by the decision function.
Implementation and an example running report using
Jupyter
Notebook are here https://anonymous.4open.science/r/sVAE-E2DB/.
A. S UPPLEMENTARY E XPERIMENT R ESULTS OF A LL D ATASETS
MIMIC-III : Each of the above-mentioned datasets only has a single prediction outcome associated to the inputs, and are not publicly available.
To extend our experiment to the scenario of multitask training on a public dataset [Table] we have conducted experiment on the clinical dataset MIMIC-III
Johnson et al.
Out of the 53,423 distinct hospital admissions for adult patients in MIMICIII, we identified 57,786 valid distinct ICU stays in total, where multiple stays might have occurred during the same admission.
Input variables include demographics, history of diseases from ICD codes
[Table]
and ICD procedure descriptions, converted to 40 numerical features after pre-processing, with missing rate approximately 11.79%.
Two prediction outcomes are extracted from MIMIC-III for all included cases: a regression task to predict the ICU length of stay (days) and a classification task to predict mortality.
This dataset is referred hereafter as ”MIMIC-III”. [Table] [Table] [Figure]
[Table] [Figure] [Table] [Figure]
[Table] [Figure] [Table] [Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure] A. L ATENT S PACE I NTERPRETATION : MIMIC-III
In MIMIC-III, we show that the model is able to learn the intrinsic information of ICD procedure codes with the promising application to billing process, as most procedures are used to treat conditions that affect a particular organ (e.g., mechanical ventilation is used to treat acute respiratory failure).
It makes sense for an ICD procedure code and an ICD diagnosis code to point to dysfunction of a particular organ.
In the 11,716 hold-out test cases, ’ventilation’, ’arteriography’, ’artery’ and ’lung’ are frequent clinical keywords, occurring in 27.71%, 14.81%, 14.60% and 4.33% cases respectively.
In latent space, dimension 5 and dimension 10 jointly represents the relevant information, and well-separate these two groups.
Meanwhile, the similarity between ’arteriography’ and ’artery’ and between ’artery’ and ’lung’ are captured, so that ICD procedures with similar meaning but different keywords are captured and clustered in all latent dimensions.
[Figure]
a: Latent dimension separating lung-disease ICU patients from artery-related ICU patients
[Figure]
b: Phenotyping of patients with artery-related ICD procedure codes
[Figure]
c: Phenotyping of patients with lung-related ICD procedure codes a & b & c: Latent dimension separating ICU patients with lung-related problems from artery-related problems
, meanwhile it captures the intrinsic characteristics within each group, although described with different ICD codes in input data (Dataset: MIMIC-III)
A. E STIMATION OF TC IN TRAINING
To illustrate the unboundedness of KL-based method in the minimization of ELBO function, we have recorded the average loss of TC term after running for 100 epochs for various latent dimensionality.
As directly estimating the TC term is intractable, we used the estimating by Minibatch Sampling proposed by Chen et al.
Two sampling strategies, Minibatch Stratified Sampling (MSS) performed indifferently from Minibatch Weighted Sampling (MWS) as shown in C2.2 in Chen et al.
(2018), hence we randomly pick MWS to record the estimated loss on training data.
For reference, we also calculated the average Pearson’s correlation coefficient between all pairs of dimensions.
As we can see, when dimensionality does larger, the estimation of TC is not lower bounded; due to the curse of dimensionality, this inevitably gets worse.
[Table]
A. M ORE C OMPARISON WITH BASELINES
We have extended Table 3 to all baseline variants of VAE models, and tabulted their classification performance (Table A4 and A5), disentanglement and reconstruction error (Table A6) below.
As we can see, our proposed sVAE continues to outperform baseline models in prediction and disentanglement, irrespective of the latent dimensionality.
In reconstruction, it is noteworthy that β -TCVAE has the best reconstruction performance than other VAE models
[Table]
[Table]
[Table]
[Table]