Surgical Prediction with Interpretable Latent Representation
Authors: Bing Xue
York Jiao
Thomas Kannampallil
Joanna Abraham
Christopher Ryan King
Bradley A Fritz
Michael Avidan
Chenyang Lu
Venue: ICLR 2022 Submitted (2022)
Keywords: machine learning
healthcare applications
latent encoding
surgical predictions

Abstract: 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.

Under review as a conference paper at ICLR 2022
SURGICAL PREDICTION WITH INTERPRETABLE LA- TENT REPRESENTATION
Anonymous authors Paper under double-blind review
ABSTRACT
Given the risks and cost of surgeries, there has been significant interest in exploit- ing predictive models to improve perioperative care. However, due 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 (s VAEE), representation learning framework for perioperative data based on variational autoencoder (VAE). sVAE provides a holistic approach com- bining salient features tailored for surgical applications. To overcome perfor- mance limitations of traditional VAE, it is prediction -aaiiiionnniddddddeeee with explicit expres- sion of predicted outcome in the latent representation. Furthermore, it disentan- gles the latent space so that it can be interpreted a clinically meaningful fashion. We apply sVAE 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 mor- tality. 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.
INTRODUCTION
Hospitals have shown significant interest in exploiting machine learning to predict surgery-related outcomes from postoperative complications (Xue et al. 202I) to surgery duration (Jiao al. 2020).
More than 10% of surgical patients experience major postoperative complications, leading in- creased mortality risk and need for a higher level of care (Tevis & Kennedy 2013). Early iden- tification of risk factors can be crucial to early intervention and improved outcomes (Kable al.] 2002). Periopertive care is also a major contributor to overall hospitalization expenses (Childers & 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 di- mensionality and complexity of perioperative data (Li et al. 2020). There are hundreds of perioper- ative 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 be interpretable from 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 ex- tract the hidden nonlinear relationship of input variables while reducing the input dimensionality through latent representations (Gao al. 2016: Kingma & Welling 2013: Rezende & Mohamed, 2013). However, traditional VAEs have two major limitations for surgical predictions. First, as VAE training is traditionally unsupervised, predictive models based VAE-encoded latent variables can suffer performance degradation (Feng et al.) 20200 Locatello et al. 2019). Furthermore, the latent space from vanilla VAE models are often entangled (Burgess et al.
2018 Chen et al. 2018), which

---

Under review as a conference paper at ICLR 2022
makes it difficult for clinicians to interpret the nonlinear encoding and trust the predictions. Recent advances in VA 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álsson et al., 2019 Cheung et al., 2014; Wu et al. 2019). but such method does not guarantee dis- entanglement 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. 20118) 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), representation learning framework specifically designed for high-dimensional and complex clinical data. s sVAE provides a holistic framework that integrates the respective strengths of regularization based disentanglement (Burgess al.. 2018: Chen al.) 2018) and guided encoding (Zhao et al. 2019 Feng al. 2020 Zhou & Wei 2020). It enhances predic- tive 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 s 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 infor- mation of original clinical data and achieves better recovery of original data.
broader setting, we applied
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-ding sigh-dimensional and noisy data.
2 RELATED WORKS
As representation learning framework, traditional VAE has been applied different clinical ap- plications such as abnormality detection (Baur et al.] 2018) or clustering tasks (Zhao et al.] (019)) Recent research advances VAE in separate directions including supe-vsee 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-V (Sorrenson et al., 2020) and others (Wu et al., 2019 Pálsson et al. 2019) organize the latent distribution with respect to the prediction (or conditionin)) 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 represen- tations should have been well-scparated 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 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 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 usually violated in model training without a designated constraint. [Hejna et al.) propose fac- tor 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 or- ganizes all samples 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 predicted values.
2

---

Under review as a conference paper at ICLR 2022
Unsupervised disentanglement. In B- VAE (Burgess 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 disentanglementt (Kim & Mnih [2018 Chen et al. 2018 Achille & Soatto 2018). As computing TC is intractable, several estimation metrics proposed (Kim & Mnih 2018 Chen 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 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
We denote X E R=Xp 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. Al. y E R represents the prediction tasks, either discrete (classification prediction) continuous (regression prediction), and is only available after the surgery-
Notations:
vanilla VAE consists of an encoder and a decoder, linked by a loss function. The encoder is usually represented by neural network with parameters Φ, which outputs (z(xx) posterior Rdd The decoder is another neural network,
Gaussian probability density of latent variable denoted by Po (x: 2). 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:
11 0) = E (=)an()(()))) [log po [2)] - K ==1
(1)
The regularizer measures the Kullback-Leibe divergence between the encoded data's 's distribution and the assumed distribution of 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) [Sor-] renson et al.] 2020), prediction label is used for conditional distribution to obtain improved disen- tanglement of latent dimensions. The loss function is defined as:
N
1(6 0)
- E-EEzag E=) (oo [log Po (x (z), u;)] + KL(gs(>/z, uz) [p(=|u;))
(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 smooth-ELBO (Feng et al. 2020). defined by:
1( 0) (x) Pe ( [z)] KL(f( (las (1)) + (Lo f/=))[((z,
where f(-) is 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 heavily penalized
3

---

Under review as a conference paper at ICLR 2022
KL divergence achieves a high degree of disentanglement in image datasets (Burgess et al. 2018 The extra scalar parameter on the KL divergence is 3 hence this variant is referred as 3- V-V is not made explicit why penalizing K [ (2 [=(()) can lead to disentangled transformations, until the ablation study the decomposition of the KL divergence showed that the total correlation is the key factor. By emphasizing on the TC term
(TC) term: TC (=)g (zz)))
with parameter /3 in the loss function, we get 3-TCVAR and FactorVAR (Chen al. 2018 Kim & Mnih 2018):
N
Σ
(a (og [L(go()/)) + 3 TC(2)
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
4 SVAE
The aforementioned works have significant implication for surgical predictions. Conceptually, dis- entanglement helps reorganize the information in X into clinical characteristies/phennoyy 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 interested in. Hence with such goals, want sVAE to 1) express explicit predictions on a latent dimension (serving as a predictor); 2) minimize the correla- tion between latent dimensions 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.
Now consider a single case x E X: as the prediction task y is dependent on (otherwise the problem is meaningful), it is also dependent fact can be represented as: p(y. x, =) p(y ())((((()))))))))))))))))))))() p()x). To distill the prediction information y from z. let g(-) be disentangling transformation function such that g(z) = [y, z'), where x' is vector of the remaining information in that is irrelevant to y. As 9(-) is disentangling function, all dimensions in z' are uncorrelated with each other too. This is similar an independent component decomposition except that we explicitly isolate y from other dimensions. Let v [v, 2²), then p(v) = (()))(='). Since v is transformation z, it is equivalently another latent representation of æc Assuming 9(-) is bijective, then p(y, x, z) can be re-written as p(y, a, z) = p(y, æ. v) p(uly, x)p(s) )p()) = p(= ((yp))p((())))
Distillation of Prediction.
its latent encoding 2. This
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 designated dimension. We neural network to implement g(z). then its generative version (z(x(())())) is equivalently another encoder. Let g(gs(2 (*)) qer (+)= now leverage
the idea in Eq.3 for the learning of p(2'(æ: and p(ylan (r),, and the TC term in Eq.5 for our assumption of disentanglement to hold true:
(6b',,a,33)
Σ (ole plog + (1)))1)))) (5)
i=1
+ KT(g (v))) ())pp(())) + B TC(=)
where a and B are the scalar coefficients for prediction loss and TC loss. Note that the smoothing function I (x)( is unnecessary, especially when y is a regressional task. In our experiments, replace it by the posterior Gaussian likelihood log qø / (1*(() (x) where (+) is the ground truth label, and set hyperparameters a = 1 and / = 10. Compared to other prediction-guides models, the introduction of TC term not only guarantees the conformation to our independence assumption, but also eliminates the use of an additional predictor prediction is explicitly expressed on single dimension.
Estimation of Correlation. Recent studies show that the estimation of TC is usually misleading. As TC is measured by KL--divergence from batch data B, TC KL(gg())) (+))) x (x)
(24)))
Jg A B) (z) log [7 (x,)) When small, such approximation is appropriate. When
4

---

Under review as a conference paper at ICLR 2022
large, there several issues. First, when sampling g0(z). due to the curse of dimension- ality, the imput sparsity leads the estimation 4e(n)()) 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 KI. diver- gence, and this inappropriaee estimation would impose negatively large, unbounded fluctuation in the loss minimization. Second, the estimation of TC using such method with mean representation is consistent with the sampled representation (Locatello al., 2019). Although several modified estimation methods were proposed (Kim & Mnih 2018 Chen et al. 2018). with higher dimen- sionality, 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 qar (#/2z) is usually smaller, amplifying such negative effects (see more results in Appendix A6).
To tackle this issue, we propose simple estimation of disentanglement using singular value de- composition (SVD). The intuition behind our method is as follows: if the latent dimensions are disentangied, linear transformation any new coordinates. In fact, it has been observed in medical imaging data that PCA/SVD-based metrics are equivalently effective 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:
new coordinate system will not increase the variance
1. Draw Monte Carlo sampling from the latent distribution of current batch data (batchsize B > d) Concatenate into B x d matrix A [V1, V ... Vnl² As each latent dimension may have different mean/variance, each dimension should be standardized.
2. Perform SVD to get the largest singular value d₁, where σ 2 02 2 ... 2 aa 2 0.
3. Calculate TC(v)= d (similar to the variance explained by the first principal compo- nent).
main advantage of this correlation estimation method is the lack of hyperparameters or expensive computation, which should provide more stable training in efficient way. Besides, proposed evaluation strategy has key properties:
1. TC(v) is bounded by [0, 1];
2. (C(v) is minimized iff the all latent dimensions are uncorrelated (orthogonal);
3. TC(v) TC(((1, (..,)...,,?)
estimated from single batch data is the lower bound of the TC of the whole latent space, calculated by TC(((1, (22) ",.,,,..,,.,, 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.
5 EXPERIMENTS
We evaluate sVAE on the open MIMIC-II 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 clinical perspective.
5.1 DATASETS
Delirium Data: The first dataset was acquired from the electronic anesthesia records of all adult pa- tients undergoing surgery at 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 binary, indicating whether delirium was developed after surgery. In total, there 12,904 patients with 52.65% of them developed postoperative delirium. This dataset referred hereafter as Delirium Data.
The second dataset acquired from large healthcare system spanning across eight hospitals, academic and community centers and surgical specialties from March 1, 2019 to Octo- ber 31 2019. The input variables include the scheduled surgery duration, clinical notes describing
OR Data:
5

---

Under review as a conference paper at ICLR 2022
Table 1: Upper: Classification performance for MIMIC-III mortality prediction. Lower: Regression performance for MIMIC-III ICU LoS. Mean and standard error are reported after 5 random shuffles.
Tanntormtionn
Direct Productio
XCrhond
SVM
DNN
Methed
Average
Average
ROC:
Average
Average
Aseraee
(d=110)
AUC
Precision
AUC
Precision
ADC
Precusion
AUC
Precision
AUC
Precuson
PCA
(X0884)
(,09)
I.A.
5085L1512)
4924
(ASTT))
7231 (9221) 579N (J052)
6488
5173
(0755)
(116)
J962
AF
3646
(A058)
VAD
200MC0X08)
3593
(0336)
725T
3556
5901 (O355)
ASRT (.007)
3-TCVAE
5728 L0216)
(,0122) (/35/ 5062
2001 (0135)
SROS
(1197)
C/0411
3630000133
.7130
(U2296) 3718
3533
Y252
(0192)
5814
2015
$550,0094
5550
(029
smooch VAL
8355 (0026) 7291 (1008))
.839
L0020)
7350
LO000)
7178
(R007)
8339
(0025)
7223 (0GA))
X337
C0025)
7200
(100
IVAE
X358 (.028) 7731 (.0048)
X367 (ON2E)
3358
(,926)
X363 (0022))
73330
Raw Dats (A=38)
(89960088)
18799
(1090)
(0004)
(.0041)
3022 (0000)
6812,
(0000)
(1000)
(0076)
Metled
Direct
(d-10)
Prediction
XGBoost
SVM
DNN
FCA
7,382 1230
RCA
(O8R6)
2585 1955000000
10:9912766
-184 (OSSE)
(2504)
(0S14)
AE
3591
(0390) 7.795
1671 (0000
1.07711351
3721702077
3.285
(1217)
(397
VAE.
1.015
2146
00331 1,851 (000
7305 (2004)
10:071 1 122054 02161 (20010400)
M.256 (1165)
8 ICYAE
9.101 (.3345)
1993
(0553) 7,803
(2831
4421 (2415)
(2.12
0108 (0032)
1.3100
.442 (.055)
P--VA
(8:901(2299))
21907
091191
8.
33555
2137 (1713)
D.0KS (1.024)
1934
(, (1751)
smooth-Al VAC
(571012363) 5618 (0307)
6.666122412)
5705
(0314) 66838
(266))
1541 (0090)
6776122300
5578 (0330)
(256)
5829
(,030)
(435)
350
6.55W (2221)
(025)
5728
(22199
Row Iara (0-38]
7,920,25455
6.514).348)
6,7500 26333
3355,01377)
6.5888
.50 7.13119
surgery type, and categorical strings (surgeon name, anesthesiologist name and location of the oper- ating 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 al. 2016) for two prediction tasks: a regression task to predict the ICU length of stay (days) and classification task to predict mortality. Due to space limitation, the supplementary results on the three dataset are reported in the Appendix A4 to A7.
5.2 BASELINES
We compared SVAE against state-of-art latent encoding and data transformation methods. For deep generative models, have included VAE, pi-VAE that organizes latent space by conditioning on prediction outcome (Zhou & Weil 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, 3-TCVAE which is an unsupervised disentangled VAE. Common en- coding/raansformation 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 lower-dimensional representation (d = 10) for all methods.
To evaluate the prediction performance using the encoded/transformeed data, several linear and non- linear 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 A3.
5.3 EVALUATION METRICS
We use the following metries to evaluate the downstream prediction performance, latent disentan- glement 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 Receiver- Operating 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². As both smooth- VAE and the proposed sVAE explicitly express the predicted outcome 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.

---

Under review as a conference paper at ICLR 2022
Latent Disentanglement: Previous studies (Kim & Mnih 2018: Chen et al. 2018 Burgess 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) assess disentanglement effects. This method relies the assumption that for particular dataset and 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 com- paring 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, essential to have 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.
0.50
Lobere D00
0.25
0.00
Latent Dim
Looore Din
-0.25
0000
0,702
1,600
Latent Dim
Actull Duraton (homal
Delirium Data
b. OR Data
Figure 1: Latent space stratification by prediction tasks in Delirium and OR datasets. In classifi- cation prediction (Fig I.a,, cases with positive postoperative delirium are shown in red and cases with negative postoperative delirium are shown in blue. In regression (Fig I.b), cases are divided into lower quartile (cyan), inter-quartile (green) and upper quartile (blue) predicted/actual surgery duration. Vertical lines represent the center of mass for each group.
bye
-05
-45
a: Encoding of albumin levels
b: abdominal surgeries from
Separating intra-
Phenotyping of intra-
(Delirium Data)
abdominal surgeries (OR Data)
other types (OR Data)
Figure 2:
Latent space separating patients by phenotypes. Circles represent the cluster of a particular group with radius equal the average distance between in-class samples and cluster center. a) Information in albumin level is split into two dimensions, as albumin is relevant to prediction. & c): Latent dimension separating intra-abdominal surgeries from other types while capturing the intrinsic similarity of intra-abdominal surgeries with different free-text descriptions.
5.4 QUANTITATIVE COMPARISONS
s VVAE has better prediction performance. When fixing latent dimensionality to = 10, we observe that SVAE produces consistently better ROC AUC and Average Precision in Delirium data

---

Under review as a conference paper at ICLR 2022
Table 2: Effects of correlation estimation on prediction, disentanglement and reconstruction of VAA with varying latent dimensions. Higher values of ROC AUC, Average Precision, and UDR are better, and lower values of Average Pearson Correlation and NLL are better.
d-10
2-15
d=20
0-25
d=330
d=552
Discnranglemest loment UDK
SAAE <VAF
(018)0405)
.165(,0227)
.1670K 0216)
(490000100)
.2491,0108 .1611,0127)
2752((0301)
16291.0110)
1409L 108
19 (11900054)
(982(,002))
(954.0006)
Prediction: ROC AUC
VAAE
.78440085) 7761000955
.7784.0089)
(7772[00067)
.7544,0091) .7691,0911) .7866.0094)
SVAE
3577L8097) 75810,080)
7576(0108)
.7585 (079)
7553(0095) (75440093)
7753110086)
Prediction: Average Precision
VAE
.945(0139) (943.0152)
7960001711
.1947(0146)
(949(,0150 (7911,01588
.7978,0148)
&VAE
(73440156)
(71860116)
3718.0171)
.771 0159) 7740,0061) (756)00146)
.7781(0147)
Recorstractoo:: Negative Log lichibood
SVAE
20560:3799 24536319)
18110,662
3165
1214:1709) 2015
3309)1133)
&VAE"
(160(.1522) (699).2611 278002240
3064033922
4910050277 3589).3884) 1.314(00125
Table 3: Comparison of disentanglement (UDR) of VAE models in three datasets. Higher values in each row imply better disentanglement.
Dannset
VAE
3-TCVAE
Pi- VAE
smooth- VAE
AVAE
MIMIC-III
.088(0162)
1164(0110)
2150(0334)
2501(,0092) 2965(00131) 1789 C0059)
.663,0158))
OR
1889101911 223800226)
311000380)
33571,0111)
Delirium
.377700077 1375(0127)
1939L0084)
.2655602277
(Table 1) and lower iMSE (higher R²) 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/disentaggling the raw data (8-TCVAE) does not benefit the downstream tasks, consistent with previous results by Locatello et al. (2019). Compared to unsupervised AE/VAFs, Vi-VAE has slightly better prediction performance, but much lower than models that introduce prediction loss into the loss function Aamooth-VAL and sVAE).
s SVAE does not require a predictor. s SVAE distills the predictive information into a single designated dimension and force the other dimensions be disentangled from this dimension. Consequently, adding a predictor has little improvement over the direct prediction provided by sVAE. This des- ignated prediction dimension visualized in Fig. 1. For Delirium Data (Fig. L.a), we force data stratification on latent dimension O 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 A4).
s VAA has consistently better disentanglement. To evaluate the disentanglement effeets of the proposed correlation estimation metric, we replaced the (C(w) by the TC term in 3-TCVAE and the Minibatch- Winibatch-Weighted Sampling proposed Chen et al. (0188), which is denoted as SVAR We compared SVAETC with s 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, full-dimensional latent space (d = 562) is also created for reference. At each dimension, we calculate UDR A2A (Duan al. 2019) for pair- wise comparison, and observed that sVAE has consistently better disentanglement than s STAE 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 8-TCVAE does not work well in disentanglement when compared to pi- VAE. Mcanwhile, 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.
s SVAE has consistent prediction and reconstruction performance with dimension reduction. Due to the bounded values of TC(x), SVAE is able minimize other loss terms, achieving bet-
8

---

Under review as a conference paper at ICLR 2022
ter distillation of predictive information and better imput reconstruction. When varying the latent dimensionalityy, the prediction and reconstruction performance is stable. On the other hand, the wrong estimation of TC term in s sVAE introduces instability in reconstruction, leading negative log-liklihood (NIL) to 0.4916 when = 25. In extreme cases (d = 562), the unbounded TC loss dominates the model training and NLL deteriorates 1.314.
5.5 INTERPRETABILITY THROUGH DISENTANGLEMENT
We now demonstrate through examples that disentanglement enhances the interpretability of the la- tent 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 characteris- tics 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 A5.
In Delirium Data, the imput 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 (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 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, esopha- gogastrnduuddeoscoppy laparoscopic, and colonoxcopy). 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, re- covery 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. the latent space, 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 surg- eries with respect 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 abdomen clustered 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 used to describe intra-abdomen regions.
CONCLUSION
In this paper, we present SVAE, a prediction-based, interpretable deep latent model. sVAE provides holistic framework that integrates regularization-ben disentanglement and prediction-uuice en- coding. We apply sVAE to two real-world perioperative datasets and evaluate its efficacy and per- formance 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- nigh-dimensional and noisy data.

---

Under review as a conference paper at ICLR 2022
7 ETHICS STATEMENT AND REPRODUCIBILITY
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 supplemtary results were attached in Appendix A4-7. Detailed description of model hyper-parameters and Python code repository can be found in Appendix A3.
REFERENCES
Alessandro Achille and Stefano Soatto. Information dropout: Learning optimal representations through noisy computation. IEEE transactions on pattern analysis and machine intelligence, 40 (12):2897-2905, 2018.
Anna Okula Basile and Marylyn DeRiggi Ritchie. Informatics and machine learning to define the phenotype. Expert review of molecular diagnostics, 18(3):219-22) 2018.
Christoph Baur, Benedikt Wiestler, Shadi Alharqouni, and Nassir Navab. Deep autoencoding models for unsupervised anomaly segmentation in brain mr images. In International MICCAI Brainlesion Workshop, pp. 161-169. Springer, 2018.
Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Des- jardins, and Alexander Lerchner. Understanding disentangling in \beta -vae. arXiv preprint (rXivv1804.035900 2018.
Ricky TQ Chen, Xuechen Li, Roger Grosse, and David Duvenaud. Isolating sources of disentangle- ment in variational autoencoders. arXiv preprint 1802.0880000494 2018
Ze Cheng, Juncheng Li, Chenxu Wang, Jixuan Gu, Hao Xu, Xinjian Li, and Florian Metze. Re- visiting factorizing aggregated posterior in learning disentangled representations, arXiv preprint (xXi::2009.05739, 2020.
Brian Cheung, Jesse A Livezey, Arjun K Bansal, and Bruno A Olshausen. Discovering hidden factors of variation in deep networks. arXiv preprint 112.v.141..6583, 2014.
Christopher P Childers and Melinda MaggardsGibbons. Understanding costs of care in the operating room. JAMA surgery, 15334)) 2018.
Brian Denton, James Viapiano, and Andrea Vogl. Optimization of surgery sequencing and schedul- ing decisions under uncertainty. Health care management science, 10(1):13-24, 2007.
Sunny Duan, Loic Matthey, Andre Saraiva, Nicholas Watters, Christopher P Burgess, Alexander Lerchner, and Irina Higgins. Unsupervised model selection for variational disentangled represen- tation learning. arXiv preprint arXiv: 12614, 2019.
Hao-Zhe Feng, Kezhi Kong, Minghao Chen, Tianye Zhang, Minfeng Zhu, and Wei Chen. Shot-vae: Semi-supervised deep generative models with label-aware elbo approximations. arXiv preprint arXiv:2011.10684, 2020.
Yuanjun Gao, Evan W Archer, Liam Paninski, and John P Cunningham. Linear dynamical neural population models through nonlinear embeddings. Advances in neural information processing systems, 29:163-171 2016.
Dean-Marie Guyader, Wyke Huizinga, Valerio Fortunati, Dirk HJ Poot, Matthijs Van Kranenburg, Jifke F Veenland. Margarethus M Paulides, Wiro Niessen, and Stefan Klein. Total correlation based groupwise image registration for quantitative mri. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition Workshops, pp. 186-193, 2016.
Joey Hejna, Ashwin Vangipuram, and Kara Liu. Improving latent representations via explicit disen- tanglement.
10

---

Under review as a conference paper at ICLR 2022
York Jiao, Anshuman Sharma, Arbi Ben Abdallah, Thomas M Maddox, and Thomas Kannampallil. Probabilistic forecasting of surgical case duration using machine learning: model development and validation. Journal of the American Medical Informaties Association, 27(12):1885-189) 2020.
Alistair EW Johnson, Tom Pollard, Lu Shen, Lehman Li-Wei, Mengling Feng. Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G Mark. Mimic-iii, a freely accessible critical care database. Scientific data, 3(1):1-9, 2016.
AK Kable, RW Gibberd, and AD Spigelman. Adverse events in surgical patients in australia. Inter- national Journal for Quality in Health Care, 14(4):269-276, 2002.
Hyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoen- coders and nonlinear ica: A unifying framework. International Conference on Artificial Intel- ligence and Statistics, pp. 2207-2217. PMLR, 2020.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In International Conference on Machine Learning, pp- 2649-2658. PMLR, 2018.
Diederik P Kingma and Max Welling.
Auto-encoding variational bayes.
arXiv preprint
arXiv: 14, 2013.
Diederik P Kingma, Shakir Mohamed, Danilo Jimenez Rezende, and Max Welling. Semi-supervised learning with deep generative models. In Advances in neural information processing systems, pp. 3581-3589, 2014.
WANG LE,, Pamela A Shaw, Hansie M Mathelier, Stephen E Kimmel, and Benjamin French. Eval- uating risk-prediction models using data from electronic health records. The annals of applied statistics, 10(1):286, 2016.
Yue Li, Pratheeksha Nair, Xing Han Lu, Zhi Wen, Yuening Wang, Amir Ardalan Kalantari Dehaghi, Yan Miao, Weiqi Liu, Tamas Ordog, Joanna M Biernacka, et al. Inferring multimodal latent topics from electronic health records. Nature communications, 11(1):1-17. 2020.
Francesco Locatello, Stefan Bauer, Mario Lucie, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. international conference on machine learning, pp. 4114-4124. PMLR, 2019.
Sveinn Pálsson, Stefano Cerri, Andrea Dittadi, and Koen Van Leemput. Semi-supervised variational autoencoder for survival prediction. In International MICCAI Brainlesion Workshop, pp. 124- 134. Springer, 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché- Buc. E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024 8035. Curran Associates, Inc., 2019.
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Rexearch, 12:2825-2825 2011
Danilo Rezende and Shakir Mohamed, Variational inference with normalizing flows. In Interna- tional conference on machine learning, pp. 1530-1538. PMLR, 2015.
Peter Sorrenson, Carsten Rother, and Ullrich Köthe. Disentanglement by nonlinear ica with general incompressible-flow networks (gin). arXiv preprint 2xXiv:2001.04872 2020.
Sarah E Tevis and Gregory D Kennedy. Postoperative complications and implications on patient- centered outcomes. journal of surgical research, (81(1):106-11 2013.

---

Under review as a conference paper at ICLR 2022
Chuhan Wu, Fangzhao Wu, Sixing Wu, Zhigang Yuan, Junxin Liu, and Yongfeng Huang. Semi- supervised dimensional sentiment analysis with variational autoencoder. Knowledge-Bass Based Sys- tems, 165:30-39, 2019.
Bing Xue, Dingwen Li, Chenyang Lu, Christopher R King, Troy Wildes, Michael S Avidan, Thomas Kannampallil, and Joanna Abraham. Use of machine learning to develop and evaluate models using preoperative and intraoperative data identify risks of postoperative complications. JAMA network open, 43)xe212240-c2
2021.
Dan-Feng Zhang, Xian Su, Zhao-Ting Meng, Fan Cui, Hong-Liang Li, Dong-Xin Wang, and Xue- Ying Li. Preoperative severe hypoalbuminemia is associated with an increased risk of postoper- ative delirium in elderly patients: results of secondary analysis. Journal of critical care, 44: 45-50, 2018.
Qingyu Zhao, Nicolas Honnorat, Ehsan Adeli, Adolf Pfefferbaum, Edith V Sullivan, and Kilian M Pohl. Variational autoencoder with truncated mixture of gaussians for functional connectivity analysis. In International Conference on Information Processing in Medical Imaging, pp. 867- 879. Springer, 2019.
Ding Zhou and Xue-Xin Wei. Learning identifiable and interpretable latent models of high- dimensional neural activity using pi-vae. arXiv preprint urXiv:2011.0440 2020.
12

---

Under review as a conference paper at ICLR 2022
A APPENDIX
A.1 ESTIMATION OF CORRELATION
Generate latent sample from v0 (z|x()))) where 1,... B are the indices of samples of current batch data. Standardize and concatenate into B x d matrix A = [0i, U2,..., Uyl" As each latent dimension may have different mean/variance, each dimension should be standardized, yielding A. Perform singular value decomposition, we get A = UEVT= = >i= diui where σ1 2 02 2 ... dd 2 Then (C(v) has 3 key properties:
1. TC(v) 7 is bounded by [0, 1];
2. TC(v) is minimized iff the all latent dimensions uncorrelated (orthogonal);
3. TC(v) (C(( 6(2) 3.., ((B) is the lower bound of the TC of the whole latent space, calculated by TC(b(), v)(2) o(N)). Meanwhile, the minimization across all batches minimizes the upper bound of the whole latent space.
A.1.1 PROPERTY 1
To see note that A is standardized therefore by definition we get |||||| E Eÿ=1 AP d. Besides, it is easy show that
|| F
o
which implies that for the largest singular value have 1s S Vd.
A.1.2 PROPERTY 2
To see 2: when (C(v) is minimized, 0 1. As is the largest singular values and by Eq. 9, UEVT UVT. Since A now product of orthogonal matrices hence itself is also orthogonal. In the reverse direction, if all latent dimensions are uncorrelated, we get AAT = 1, so the singular values of A are all and TC(v) = =0.
we get 01 = 02 = ad 1. Therefore,
For the property 2, intuitively TC(v) measures the maximum variance that can be explained by principal component. When the largest principle component can only explain minimum variance, know that the latent space is decorrelated (disentangled).
A.1.3 PROPERTY 3
To prove 3, we first show that TC(v) is the lower bound. Note that the singular values of A are square roots the cigenvalues of A A. When estimating the whole latent space (B N), this is equivalent to
(AT
X
A X
= A"
where X are the latent variables of data not in the current batch (X contains N - columns). So the question is equivalent estimating the eigenvalues of the sum of symmetric (in this case also positive semidefinite) matrices using only A From Courant- min-max theorem, we know that the eigenvalues of AT A are:
B rows and d
A (A= A) min mmax RA (x))2 U and æ + 0}} dim(U) =k}
where RAT (x)) Ao. (a,z)
Since R xxxx+ 2 max max(RATA) A Rx (x)), we know that (ATAT) S (ATA + X X), hence (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,, and then the whole data can be expressed as
A
D =
(8)
x
13

---

Under review as a conference paper at ICLR 2022
we know that =²D== + +X7X, X, Then
TC(D) - x
hence
DI
A
2 1 d
A
X 112 2-1
TC(D)
d-1
||| A 122 1
A 112
X2X, 1
d-1
d-1
d-1
(T) +
After a complete epoch looping through all training data, we have actually minimized the upper- bound the the true estimation. When the training converges, would know that max(TC((v)] - (C(D)= I (v))
14

---

Under review as a conference paper at ICLR 2022
A.2 GRAPHICAL DIAGRAM OF SVAE
As raw EHR data consists of non-numerical or non-static variables, standardized data-processing was conducted for each dataset get the static, numerical representation of X illustrated in Fig.
A1.
Clinical Texts Teatual
Frebsdding
Categorical Data
Factorization
Encoder+ + Transformation
P po (1
Decoder
X"
Numerical Variables
Time Series
Feature
Outcome y
Extraction
Figure Al: Graphical diagram of sVAE. Heterogeneous mixed type data are first converted to merical embedding, and then fed into the encoder. Prediction task is only visible by sVAE during the training stage to regulate the latent space, and is not needed when encoding the testing data
15

---

Under review as a conference paper at ICLR 2022
A.3 HYPER-PARAMETER PARAMETERS AND IMPLEMENTATION DETAILS
All deep latent models, including AE, VAE, SCCVAE.smooth-VAl and SVAE are implemented using Python Pytorch (Paskke 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= 1e ). batchsize 256, number hidden layers 4 (same for both encoder and decoder), and number of hidden units (10%d, h. h h) for decoder and (h, 8*,, 4 h) for encoder, where h is ayper-parameter Currently h 2000. The DNN predictor 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 imple- mented by Sklearn package Pedregosa et al. (201T). For LR, we used 11 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 "Is" in regression and "binary.logistic" in classification. LogisticRe- gression(penalty=111 sool=vr "solver='iiblina ,C=0.2). For SVM, the ROC AUC and average precision are calculated by the decision function.
Implementation and an example running report https//aaoyyymouu...aaamassiiiis
using Jupyter Notebook are here
16

---

Under review as a conference paper at ICLR 2022
A.4 SUPPLEMENTARY EXPERIMENT RESULTS OF ALL DATASETS
Each of the above-mentioned datasets only has single prediction outcome associated to the inputs, and are not publicly available. To extend our experiment the scenario of multi- task training on a public dataset, we have conducted experiment on the clinical dataset MIMIC-III Johnson et al. (2016). Out of the 53,423 distinct hospital admissions for adult patients in MIMIC- III. 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, 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: 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".
MIMIC-III:
Table A1: Classification performance for postoperative delirium prediction.
Frastormation
Direct Pediition
LR
XCHooost
SVXI
DNN
ROC
Average
RCC
Average
RDC
Average
ROC
Average
ROC
Averge
AUC
Precision
AUC
Precision
AUC
Precisioo
ALC
Precision
AUC
Precision
RCA
7168
7594 (0151) 7002 (002)) 7254
62997
117T00270 (0213))
452 (17633)
7221 (OMI54)
(0155)
ICA
7468 (089)
7695
(0122)
(066) 7013 (01231
6878 (0016) 2048 (0090)
7455 C00R6)
7665 (,0277
(311000130)
6463)00066
(0090)
(0099)
6677 10107)
(0089)
.6400 .0117)
Д-ТCVАE
(363020114)
6722 (0071)
1333
6723
1
(0692)
6712
(0148)
(062;0016)
pi-VAE
X900X
7135)0183)
66AAN
0108)
681410155)
6871
(0173)
(0215)
6684
(0)
1688
smooth VAE
J388 (.01U5))
27722 GOLRT)
775L LOL
7584
7723
7508
(M095)
7773 (0135)
7558
LI0E2))
77722,001300
sVAB
3761 (1995)) (8152)
77225
(.0101)
.174).0144)
.7.35
(.091)
(781300155)
7691 (.097)) 7862100631
7746
(J099))
.7966.0161)
Data (6-562)
.757% (091)
(0153)
2374 (O100) 7660
(1127))
7459 (007)) 2001 (0149)
(11111) .7629 (00146:)
Table A2: Regression performance for predicting length of stay in operating room.
Method
Direct Prediction
(d-10
LR
XGBoast
SVM
DNN
IMSC
IMSE
R
IMSC
IMSE
PMSE
R
PCA
94.33(29155)
A52000520)
(3.3303014)
78.1713.993)
6429.0025
82.5022174)
(027.0275)
KCA
111.4
2717.1990
2508
(8838)
(462(0083)
AE
4799 0225)
35 65
5589
0196001991
VAB
101,
(3,10)
A050
(003)
10 722985
13810075
16 16..9..200
344% (4005
3: 3UVVL
100,415:22
55 (8,422
2017000210
100.3
10..0 16.393
3199405544
pi VAB
42844
92.7
(5.875)
4965
(6-48)
79.
6285
0722
K5 32 (5.33
5948
(0036))
Moodi-VVE
(283418857) 72322 (0226)
(772.793)
72386 (1222)
W8.99 (30665)
7293
(0344)
25.25 (2:749)
1219
7217,1997
<VAE
67.336189) 735100249)
67.1963199)
7268,0251)
98.3,000
7220
(244)
738N
0248)
68 (8.2674.667)
7273,0328)
Naw Data (6-23)
(500))
(6
36 (2665)
7251
2227 (531)
19.11 22,55)
7179 (179)
Table A3: Regression performance for MIMIC-III ICU LeCCULLhooffSSSte prediction. Mean and stan-
dard error are reported after 5 random shuffles
Traisfoemation
Method
Direct Prediction
(d-10)
LR
XGBoost
SVM
DNN
tMSE
1
IMSE
1
IMSB
R
IMSE
MSE
R
SVALTE
6.719 (511))
6.86
2560)
367900000
6615
5YR0
0390 6 7801 2550
3556 10333) 3729 6,788 (10)) SM00 0274)
(.6NT.1300
3756 (1228)
6.351
(1235))
(,0296)
[6333
(,156)
3916T
(.436)
Raw Data (2-18
79200
3954.1997
(331,,866
SWALARE
6780 2855)
3885184471
6.6AN
2465)
3517
Table A4: Classification performance for MIMIC-III mortality prediction. Mean and standard error reported after 5 random shuffles.
Tranefommaion
Direct Prediction
XGBoost
SVM Averige
DNN
Method
RDC
Averae
ROC
Asarage
Average
RCC
(r::10)
AUC
Precision
Precision
Precision
AUC
#VAE"" <VAF
85537
72991 0155)
85361 (025)
7296, (MS5))
X355L
71 0055)
MS57 00025)
7297) (005)
X355, (114))
72544 (005)
4365100021)
.7833.0046)
8372 (020)
7345(10554)
8359L0022)
.7208K 09.9))
837W.0017)
.743.0043)
(372
.7345/ (654)
Raw Dath
(0133)
A739 (1899)
3048 (1KE44)
2214 (064)
3022 (O0AR))
ART2 (TR00))
X358 (10037))
7.231 (,0066
17

---

Under review as a conference paper at ICLR 2022
Labore Dim
02
Labont Dm
0.000
0.50
1,000
Latent space organization by predicted risk of mortality (Dataset: MIMIC-III
Latent space organization by pre- dicted ICU length of stay (Dataset: MIMIC-III
Figure A2: Latent space stratification by prediction task in multi-task training. Positive cases in mortality prediction are shown in red and negative cases in blue. For regression prediction of ICU LoS, cases are divided into lower quartile (cyan), inter-quartile (green) and upper quartile (blue). Vertical lines represent the center of mass for each group.
18

---

Under review as a conference paper at ICLR 2022
A.5 LATENT SPACE INTERPRETATION 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 billing process, most procedures used treat con- ditions 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 dysfunc- tion 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% %6 14.60% and 4.33% cases respectively. In latent space, dimension 5 and dimension 10 jointly represents the relevant infor- mation, and well-separate these two groups. Meanwhile, the similarity between 'arteriography' and "artery' and between 'artery' and 'lung' are captured, that ICD procedures with similar meaning but different keywords are captured and clustered in all latent dimensions.
(ved Malen
at Latent dimension separat- ing lung-disease ICU patients from artery-related ICU pa-
b: Phenotyping of patients with artery-related ICD pro- cedure codes
c: Phenotyping of patients with lung-related ICD proce- dure codes
tients
& & Latent dimension separating ICU patients with lung-related prob- lems from artery-related problems, meanwhile it captures the intrinsic charac- teristics within each group, although described with different ICD codes in input data (Dataset: MIMIC-III
Figure A3:
Latent space separating patients by phenotypes. Circles represent the cluster of partic- ular group with radius equal to the average distance between in-class samples and cluster center.
19

---

Under review as a conference paper at ICLR 2022
A.6 ESTIMATION 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 intractable, we used the estimating by Minibatch Sampling proposed by [Chen et al. (2018). 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 '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 the curse of dimensionality, this inevitably gets worse.
Table A5: Unboundness of TC estimation (log) in training
Dimensionality
TC (MSS)
Avg. Pearson Correlation
5
-15.68 (.0513)
.0349(0069))
10
-35.24 (.1103)
.0253(.0246)
15
-54.69(.1761)
.0406(.0211)
20
-74.19(.2299
.0367(.0309)
25
-93.35(.6048)
(0944(.1036)
30
-112.7 (.4669)
.0729 (.0360)
562
-2175(7.208)
.0643(.0182)
20

---

Under review as a conference paper at ICLR 2022
A.7 MORE COMPARISON 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 disentangle- ment, irrespective of the latent dimensionality. In reconstruction, it noteworthy that 3-TCVAE has the best reconstruction performance than other VAE models.
Table A6: ROC AUC of baseline methods with respect to dimension.
d-3
d-10
d-20
2-562
Pi-VAE
6920X
(22210255)
721300103)
9079
GBT
0623 01200
A6-PL.112)
6739
.1996)0001)
0912131533 .19350077)
724440000
SVM
66199 0259)
66764 (1787)
(192)
(049(1084)
7173,0123)
01009
79774
67300019
6868C01199
6855501223
7583,00653
3-TCVAR
LR
636K01991
6567
6659(1911
6743CD161
678760142)
71-11 681200102
GBT
6115000611
0177
065100130
(010)
56500 0098)
00 (7185
SVM
6373 0144
6563..0 01003
(88100183)
61355,1500
(1161
6874
74044 (1022)
DNN
615200106]
61677
6559(0118)
6621(0146)
667100182)
6774) 0083)
73627 (0053)
Direction Prediction
0155
76088_0106
756210082)
7587(0005)
7544(0091)
75668,000
(579900079)
L
7179
(870L0088))
7559 (189)
758RK NOR
2557
754700087)
SVM
757N, 0121)
7615681022
7567 (0077)
2605 (O06))
7557 (104))
7582) (0055)
75.49 (1122)
DNN
7588 0118)
761100103
75561 (0083)
761361886
7580000877
7803,
7589I (09T)
Table A7: Average Precision of baseline methods with respect to dimension.
de3
2=10
4=13
4-20
am
VAE
I.R
MMSO 0165)
71552
(195(0050)
2378.10667)
0111)
7845 .01
GET
6781)00F)
0167)0163
6890(0995
69SAC1OSS)
NISU (1ST)
74681 01
SVM
0785501931
(22500227)
20155
(410
77466 0144
DNN
6787,0143)
6901 (226)
69288 (2241)
2012 (043)
7214 (1158)
31688 (216)
7802 2014
ATTCVAR
LR
6571C00567
6738800700
69990
699700171)
701000136)
739700123]
GBT
6619000006
6735(00061
6833)0074)
690550113) 6898(0071)
0942)0008)
69011.0067)
SVM
6560 (0229)
0090
(0391)
2225)0263)
0109)00281 705100145)
760er 01:41
6429
650PL.112
67511.0126
0883 0225)
0985(0013)
(576L0173)
VAE
Direction Prodiction
01666
778750155)
7759(014))
3765
072300145)
776200150
0763)00411
LK
1803001391
(150L01501
1791000166
775801188)
.779600150
7775501333
GBT
7688 01455
7143
7O98L 0155))
0130
TRI9 773960179)
0103)
.7733,01 212
SVM
0153,
0143
T748 0165) 2784,D133)
7766, 0161)
774UK (1149)
DNN
777200153)
7803001555
3773(0133 (7796(0131)
7776(0163)
7769100155 7800000166
Table A8: Effects of correlation estimation on prediction, disentanglemen and reconstruction of sVAE with varying latent dimensions. Higher values of ROC AUC, Average Precision, and UDR are better, and lower values of Average Pearson Correlation and NLL are better.
5=
d=10
d=15
d=20
d=25
d=30
Disenianglemenn
UDR
PA-VAE
29377,2111
19399.0044)
1372(0204) 166 (0S8) J088)
(00350055) 09
B-TVVAE
1987,0189) 1375(00127)
1138,
08 79779, (1068)
smoothhVAt
2726)0222)
.1899.0059)
1585(00033)
.1770,0050)
.171(,0036)
.099100026)
Reconstruction: Negative
Log-likelhhod
P--VAE
5896,8472)
4037(4903)
2962(+411)
80K2]
3936(,9701
4261[6888]
4119(
9-TCVA
135660242) .1684(1364)
2399(29022)
.1898 2258)
0802(0553)
0991(.11-2)
1.3532,2422
smooth- VAE
3535(,2576
,633((1380
4754(3933)
(315([3333]
2306(,
2237(,3054)
6271066844
21

---

Under review as a conference paper at ICLR 2022
A.8 GRID SEARCH OF HYPER PARAMETERS
Table A9: Classification performance when varying hyper parameters (MIMIC-III mortality predic-
tion). wt weight of prediction loss; B: weight of disentanglement.
Hyper Paranelers
ROC Avenge
ROC
Avenge
ROC Asergg
ROC Avergge
ROC
Avenge
ROC
Avenge
ROC
Averge
ROC
Avenge
(wr,2)
AUC
Precision
AUC
Precision
AUC Precision
AUC Precision
AUC Precision
ADC
Precision
AUC Precision
AUC
Precision
7210
8381
.7197
8399
-7218
8406
7215
.7274
8378
8391
.7253
8387
7218
8391
7217
-754
8403
7250
8406
7228
7782
3385
7178
8399
.2218
J=1 7
K183
7179
837
2218
X,00
7150
8356
7237
8391
7.56
3410
72TV
8378
720
8389
7197
== 10
8350
7214
8377
.173
8103
7216
8382
2210
.391
7266
8110
7219
8373
7209
8389
.713
22


Reviews:
[{"review_id": "JCe0Kw5pDg3", "reviewer": "Reviewer_CEnW", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The proposed method combines the idea of disentangled latent variables and identifying meaning of representation for the prediction task. It also provides a new estimation for the correlation between latent variables, which also showed superiority in the conducted experiments. There is novelty in the methodology, however, the marginal improvement of the proposed methods over the previously published methods may not be sufficient to justify its superiority.\nBesides that, I may have the following questions:\n1. The authors also suggested a new measure of correlations between the latent variables and  to use that instead of TC. The new correlation metric was compared with TC in the delirium dataset. Was this consistently observed in other datasets?\n2. There are two hyper parameters introduced in this method, but the authors did not mention the reasons for the choice of the hyper parameters. When we think about an end-to-end modelling process, the hyperparameter tuning and model selection process might also play a role in the evaluation of the method. Will the authors help to explain the choice of the hyperparameters?\n", "summary_of_the_paper": "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.\nThe 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. \n", "summary_of_the_review": "Compared with the previously published methods, however, the new method did not show a significant superiority. However, it was a good application of VAE for an important clinical question. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "G7HM45GiIh7", "reviewer": "Reviewer_56Ru", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "In addition to comments made in the Summary of the Paper section, some experimental issues exist. How the (severe) missing rate in the delerium data was handled is not expressed. It is good to include MIMIC-III, and the baselines are well chosen (although how they were optimized is also not expressed). Instead of choosing a single standard for all baselines, it would have been preferable to understand their strengths and weaknesses at different hyper-parameters (e.g., does one’s performance grow more or less quickly with more (or fewer) reduced features $d$? I.e., Table 3 but with all baselines). Gaussians are used w/o consideration, but this is as usual so it is not taken as a negative. Although sVAE clearly outperforms many baselines, the degree to which it’s close to smooth-VAE, especially considering (supposed?) variances would make a significance test, and some cross-validation, advisable.\n\nThe interpretability angle is interesting and appreciated, especially within the medical context. Not enough has been shown to warrant the mention of ‘trust’ in the conclusion, though. \n\nThere are a few textual issues, including grammar (e.g., “… we want our sVAE to … minimizes …”) but in general the text is well structured and constructed.", "summary_of_the_paper": "This paper proposes a representation learning framework for high-dimensional, complex clinical data. It is not clear to what extent this is specific to *surgery*, however. The related works section is appropriately structured, deep, and recent. Some of the text (e.g., around background) is needlessly simplistic. \n\nThe primary methodological contribution orbits around the function $g()$— the  disentangling function applied to the latent representation of an encoder. Naturally, there are many papers that disentangle latent representations (e.g., Yan et al 2016 (Attribute2image), Higgins et al 2017 (beta-VAE), Mathieu et al 2019), including various attempts at sequential encoding; some of this work is cited in Sec 2, and the efficient computation of the total correlation is presented as the key step to add and then overcome. The result is Eq 6 and SVD is subsequently used in the estimation/replacement of TC. While the rationale provided is _fine_, some claims made regarding the intuition for this step do not seem to be matched by the rationale/citations given, and it would have been advisable to dig into this substation theoretically, a bit more.", "summary_of_the_review": "Generally, there are no major concerns with regards to the empirical setup (nothing that was done appears incorrect or suggests incorrect results), except that more should have been done to validate the results. The technical contribution over the vast amount of other work on disentangling latent representations is clear but perhaps slight and therefore the applied angle seems to be emphasized (indeed, ‘surgery’ is even in the name of the model) but, apart from using some relevant datasets and showing some qualitative results, there is almost nothing about the applied context that seems to influence the design, nor the analysis of the results beyond the standard ML metrics.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "2OQx3-59Vy", "reviewer": "Reviewer_4Kio", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1) Regarding reproducibility, an experiment using a public dataset is more important. Thus it would be better to have MIMIC-III dataset results in the main text rather than one of the private datasets.\n\n2) While the authors argue that this sVAE is specifically designed for high-dimensional and complex clinical data, there are many prediction tasks in various fields using such high-dimensional and complex, (and noisy) data. Thus the authors should at least discuss generalization of this approach.\n", "summary_of_the_paper": "To improve accuracy of prediction and interpretability on a task using high dimensional and noisy data, the authors proposed surgical VAE (sVAE). sVAE is prediction-guided with explicit expression of predicted outcome in the latent representation, and it disentangles the latent space so that it can be more interpretable. sVAE was evaluated using two private datasets and one public dataset, and showed its efficacy on both classification tasks and regression tasks.\n", "summary_of_the_review": "The authors introduced sVAE, a prediction-based, interpretable deep latent model designed for clinical prediction tasks. However, such tasks based on high-dimensional and complex data exist not only in surgery but also in a wide range of fields in the real world, thus generalization is important considering the backgrounds of ICML participants. If the authors aim to keep the scope within the clinical or surgical field, a more specific venue would be appropriate.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviews are of good quality. The responses by the authors are commendable, but ICLR is selective and reviewers still believe that the research would be better as two separate papers: one about the problem and solution from an ML perspective, and the other about the application to surgery. Papers that provide a new method in the context of a single application domain run the risk of making a contribution to neither, and of being evaluated by reviewers who are not experts in both.
Decision: Reject