Unbiased Representation of Electronic Health Records for Patient Outcome Prediction
Authors: Lucas Jing Liu
Victor Ortiz
Javier A Neyra
Jin Chen
Venue: Submitted to ICLR 2023 (2023)
Keywords: Deep Learning
Electronic Health Records Representation Learning
Healthcare AI
Model Fairness

Abstract: Fairness is one of the newly emerging focuses for building trustworthy artificial intelligence (AI) models. One of the reasons resulting in an unfair model is the algorithm bias towards different groups of samples. A biased model may benefit certain groups but disfavor others. As a result, leaving the fairness problem unresolved might have a significant negative impact, especially in the context of healthcare applications. Integrating both domain-specific and domain-invariant representations, we propose a masked triple attention transformer encoder (MTATE) to learn unbiased and fair data representations of different subpopulations. Specifically, MTATE includes multiple domain classifiers and uses three attention mechanisms to effectively learn the representations of diverse subpopulations. 
In the experiment on real-world healthcare data,  MTATE performed the best among the compared models regarding overall performance and fairness.

Under review as a conference paper at ICLR 2023
UNBIASED REPRESENTATION OF ELECTRONIC HEALTH RECORDS FOR PATIENT OUTCOME PRE- DICTION
Anonymous authors Paper under double-blind review
ABSTRACT
Fairness is one of the newly emerging focuses for building trustworthy artificial intelligence (AI) models. One of the reasons resulting in an unfair model is the algorithm bias towards different groups of samples. biased model may benefit certain groups but disfavor others. As result, leaving the fairness problem unre- solved might have significant negative impact, especially in the context of health- care applications. Integrating both Nomain-speciinn and domain-invariant repre- sentations, we propose a masked triple attention transformer encoder (MTATE) to learn unbiased and fair data representations of different subpopulations. Specifi- cally, MTATE includes multiple domain classifiers and uses three attention mech- anisms effectively learn the representations of diverse subpopulations. In the experiment on real-world healthcare data, MTATE performed the best among the compared models regarding overall performance and fairness.
INTRODUCTION
Electronic Health Record (EHR) based clinical risk prediction using temporal machine learning (ML) and deep learning (DL) models benefits clinicians for providing precise and timely interven- tions to high-risk patients and better-allocating hospital resources (Xiao et al. 2018 Shamout et al.) 2020). Nevertheless, a long-standing issue that hinders ML and DL model deployment is the con- cern about model fairness (Gianfrancescoo al.) 2018 Ahmad et al., 2020). Fairness in AI/DI. refers to a model's ability to make a prediction or decision without any bias against any individual or group (Mehrabi et al. 202T). The behaviors of biased model often result in two facets: it performs sig- nificantly better in certain populations than the others (Parikh et al. 2019), and it makes inequities decisions towards different groups (Panch et al. 2019). Clinical decision-making based upon bi- ased predictions may cause delayed treatment plans for patients in minority groups or misspend healthcare resources where treatment is unnecessary (Gerke al. (020)
The data distribution shift problem across different domains is one of the major reasons a model could be biased (Adragna al. 2020 To address the fairness issue, domain adaptation methods have been developed. The main idea is to learn invariant hidden features across different domains, such that model would perform similarly no matter which domain the cases belong. Pioneer domain adaptation models, including DANN (Ganin et al. 2016), VARADA (Purushotham et al.) 2017). and VRNN (Chung et al. 2015), learn invariant hidden features by adding domain classifier and using gradient reversal layer to maximize the domain classifier's 's loss. In return, the learned hidden features are indifferent across domains. Recent work MS-ADS (Khoshnevisan Chil 202T) has shown robust performance across minority racial groups by maximizing the distance between the globally-shared presentations with individual local representations of every domain, which ef- fectively consolidates the invariant globally-shared representations across domains. However, it is difficult to align large domain shifts and model complex domain shifts across multiple overlapped domains.
Alternatively, the data distribution shift problem could addressed using domain-specnt -speciff bias cor- rection approaches. A recent study showed that features strongly associated with the outcome of interest could be cubpopulation-speciie together all features from patients with different backgrounds might bury unique domain-specific
(Chouldechova Roth 2018). It indicates that lumping

---

Under review as a conference paper at ICLR 2023
EHR data
HR
Domall-Specfffc Feature Relevance
Attention
[Domain-nvarian
688
Domam
Domsin Clessiferr
pattent
Domain- Specifii)
Temporal Attention
Domain
Tone
Domaann-oooaae Representaton
Outsome Prodiction
Input
Dynamic Weighting
Rx
Overall framework of masked triple attention transformer encoder (MTATE). HR, SBP, and sCr stand for heart rate, systolic blood pressure, and serum creatinine, respectively. xy represents all clinical features at time f, fi represents values feature at all time points. Z represents the data representations learned from the it feature relevance attention module.
Figure 1:
characteristics. Afrose et al. proposed to use double prioritized bias correction to train multiple candidate models for different demographic groups (Afrose et al. 2022). Similarly, AC-TPC & Van Der Schaar (2020) and CAMELOT et al. (2022) used clustering algorithms to generate representations of patients with similar backgrounds and use cluster-specifie representations for the outcome prediction.
In summary, both domain adaptation and domain-spectie bias correction approaches address the same fairness issue with different assumptions about the relationships between latent representation and the prediction outcome. The former believes that performance variation across domains would be benefited from invariant feature representation. while the latter affirms domain specific represen- and domain--pecific data representation should
tations. remains unclear whether be used for prediction task.
To better address the fairness issue, we propose an adaptive multi-task learning algorithm, called MTATE (i.e. Masked Triple Attention Transformer Encoder), to automatically learn and select the optimal and fair data representations instead of explicitly choosing domain adaptation or domain- specific bias correction. Under this setting, both invariant and domain-specific representations are special cases where one the approaches dominates the data representation. The purpose of MTATE is to generate multiple masked representations of the same data that attended by both time-wise attention and multiple feature- wise attentions parallel, where each masked representa- tion corresponds to specific domain classification task. For example, breaks the patient cohort into subpopulations defined by race, and another classifier is focused on gender. The learned EHR representations could be domain-specific, Nomain-invariant, or the mix of the two reflected by the domain classification loss values. A low loss value indicates the representa- tion domain-specific, and a high loss value indicates Nomain-invariant. The model will compute the representation wise attention for each individual testing case, leading to personalized data repre- sentation for downstream predictive tasks. The overall framework of MTATE is shown in Figure [ The primary goal of MTATE is to learn an unbiased representation to make fair and precise patient outcome predictions real-world healthcare setting.
of the domain classifiers
To demonstrate the effectiveness of MTATE, we focus rolling mortality predictions for patients with Acute Kidney Injury requiring Dialysis (AKI-D), a severe complication for critically ill pa- tients, with high in-hospital mortality rate (Lee & Son 2020). The clinical risk classification for AKI-D patients is challenging due to complex subphenotypes and treatment exposures (Neyra & Nadkarni 2021 Vaara et al., 2022). There is an urgent need to develop actionable approaches to account for patients' backgrounds and subpopulations for personalized medicine and improve patients-centered outcomes (Chang et al. 2022)
2

---

Under review as a conference paper at ICLR 2023
Domain Classifier
Gender
TR-Attention
Domain Specifiic
Poslional
FR-Attantion
Haas
Mutt-Hoad
FRyAdlendan
Suppe
Inpel
FReAllention
FR-Atlenion
Representation
E
Figure 2: Network structure of the masked triple attention transformer (MTATE) algorithm. The TR-attention and domain specific FR-Attention module can be stacked N times
The contributions of this work are three-fold: 1) To the best of our knowledge, MTATE is the first model that integrates both -ppaif- and -onainninvariant features one model. It trains the fair representation and predicts the downstream task altogether, which is critical for reliable clinical outcome predictions for patients with different demographics and clinical backgrounds; 2) MTATE employs time-wise, feature-wise, and representation-wise attention mechanisms to compose data representations for downstream prediction tasks dynamically; and 3) MTATE effectively mitigated the bias towards different subpopulations in the rolling mortality prediction tasks on AKI-D patients and achieved the best performance compared to baselines.
METHOD
MTATE consists of five components, and the detailed architecture is shown in Figure 2 The first component temporal-relevance attention (TR-Attention) module to generate time-wise attention that associates each time step with the other time steps considering all features. The output is time-attended representation. The second component is domain-specific feature-relevance atten- tion module to generate feature-wise attentions that associate each feature with the other features considering all the time steps. The outputs are multiple feature-attended representations, one for each domain. The third component is module consists multiple domain classifiers, and each clas- sifier classifies each feature-attended representation into a predefined domain. The fourth module unified data representation module, which uses the representation-ww wise attention to aggregate representations (domain-invariant or domain-specific) to final representation. The last module is outcome prediction module where the final representation is used for patient out- come prediction.
2.1 NOTATIONS
A patient's EHR data can be represented X {x1, X2. Xty XN. X R.N,,,, where N, is the number of features and N₂ is the number of time steps. X2 E R RlxN, represents a vector of clinical parameters (e.g., heart rate, blood pressure, etc.) time step We consider binary outcome and domain classification problem in this study. The patient domain class labels are denoted as dy E VR, where N₂ represents the total number of domains, dy (0, 1} represents the label for the i-th domain, 1 and 0 represent whether given patient falls in the target domain or not, respectively. The patient outcome label is denoted as y E {0, 1}. where and represent death and alive before hospital discharge.
2.2 TEMPORAL RELEVANCE ATTENTION
The temporal-relever relevance attention (TR-Attention) module aims to learn the relationships between each time step other time steps considering all features each time step. We first the position
3

---

Under review as a conference paper at ICLR 2023
encoding from the original Transformer encode the relative position information to the input X and use the multi heati-heeddde attention mechanism to learn the temporal-relevance attention.
Specifically, query, key, value vectors (Q, K. V) are the linear projections of all features at every time steps X. Thus, the attention weights computed from the query and key represent how much focus all features one single time step is associated with themselves at other time steps. Then, the output of each head Zh is the multiplication of value vectors and time-wise attention TR, The final output of the TR-Attention module Z R RAX,, is the linear transformation of the concatenation of the output of every head. Lastly, the residual connection and and layer normalization are applied Z, denoted as Z Layer NermerNNrm(2 (Z X). The temporal-relevant attention of each head is represented as:
Q. V XXQ; XWK, XWv
QKT Vlx
TR = softmard
(2)
2 A V
(3)
Z Concat Z2, ,... Zb,,... .... NN,,Wo.
(4)
For simplicity, we assume all projection matrices WQ. WK. Wy have the same dimension de- Thus, Wo, WK, Wv E xd,, Q K, V Nxxd, the temporal relevance attention is TR RNN,, the output of each head is 2h E R.xx, The projection matrix for the final output Wo (NN,, where N₁ represents the number of heads.
2.3 DOMAIN-SPECIFIC FEATURE RELEVANCE ATTENTION
The domain-specing feature relevance attention module aims learn each domain's diverse and unique latent representation. The module includes a set of parallel sub-modules called feature- relevance attention (FR-Attention), where each FR-Attention module focuses on the representation of one specific domain. Since features for the different domains are not equally important, we randomly masked out some number latent features along all time steps differently for each FR- Attention module. In return, the masking procedure forces each sub-module to Iearn differrnt feature focuses for each domain to generate unique domain representations. Then feature-wist attention is computed using the multi-head attention similar to the TR-Attention model. FR-Attention module learss the relationships between each latent feature and the others considering all time steps.
The input of each FR-Attention sub-module is ZT E NNNNN, which is the transposed output of the TR-Attention module Z. Then, Z is passed through masking layer, in which MR x Ny number of latent features are randomly selected and removed from Z where M represents the masking rate. We denote the masked input of each sub-module
M E RNN, where N₁ represents the number of features after masking. M is passed through the multi-head attention block as well as the residual connection and layer normalization. Finally, M is transposed back to the original form. Then it is passed through point- wise feed-forward network (FFN) well as the residual connection and layer normalization to get the final output, denoted as Z N,x,N, Please see the detailed architecture and formula for FR-Attention in Appendix Section A.1
2.4 DOMAIN CLASSIFIER
Multiple domain classifiers are used to classify patients into subpopulations based on the latent representation from FR-Attention sub-module. The input of one domain classifier Z E N.N, where i denotes the index of the sub-module or domain. Z is flattened by taking the max along the dimendimnnioon, then followed by a linear layer with a sigmoid function for binary classification. We binary cross-entropy for the domain classification loss, denoted as Ld,- The domain classifier module assists to learn latent representation for each domain and the domain classification loss is used for generating representation-wine attention in the next module.
4

---

Under review as a conference paper at ICLR 2023
2.5 DOMAIN-FOCUSED REPRESENTATION AND REPRESENTATION WISE ATTENTION
While a domain-specific representation module is focused on representing its target domain, the resulting representation from each domain be domain-specifie or domain-invariant according to the domain loss. However, not all representations are equally crucial to outcome prediction. Thus, this module aims to generate the final representation for the outcome prediction considering both Noomain-speciffe and domain-invariant representations and their corresponding domain loss. We call this module a domain- focusee representation module since both domain-specific and domain- invariant are candidate representations.
The input to this module is transformed version of the latent representation generated from each FR-Attention module Z, Every Z gets transformed to its original dimension by adding the masked(removed) feature back with values of so that all latent representations can be aligned in feature space. We denoted this particular form of latent representation as Zí Z is flattened by tak- ing the max along the time dimension, and all flattened vectors are concatenated together, denoted as E IR Na,N, The final representation C R is the weighted sum of all the candidate rep- attention (RW-Attention),
resentations, where the weights a E R Nax1 are the representation-we which is computed based on E and the domain prediction loss La R Nax1 following:
(5)
Na
Cj
a. B
(6)
where UA E 1R( and WA Rd,x1 are the projection matrices, i represents the domain index, j represents the feature index.
2.6 PATIENT OUTCOME PREDICTION
The final representation C of EHR is concatenated with all static features, such as demographics and comorbidity. followed by a linear layer with a sigmoid function for the outcome binary prediction. Let the patient outcome label be y and the predicted label be ÿ, and we use the binary cross entropy as the part of the final loss, denoted as Lp. We also constructed the supervised contrastive loss Khosla al. (2020) to mitigate further the model bias as another part of the final loss. The contrastive loss denoted as Le- The final prediction loss L is:
Lp + Le
N,
Lp
(y, log(()) + +1 yi) log(1 (.))
(8)
i=1
N
exp(h; hp/T)
L
log
N. E exp(hy ha/T)
where No.No. Na represents the number of all samples, the number of samples having the same labels as the anchor samples (j), and the number of samples having the opposite label to the anchor samples (j). h represents the concatenation of the learned representation C and the static features. r is a scale parameter.
EXPERIMENTS SETTINGS
In the experiment, we aim to continuously predict AKI-D patients' mortality risk in their dialy- sis/renal replacement therapy (RRT) duration. More specifically, given a period of EHR in dialysis duration before time T, will continuously predict the mortality risk between T and T + 72h.
3.1 EXPERIMENT DATA
The study population consists 570 AKI-D adult patients admitted to ICU at the University Hos- pital from January 2009 to October 2019. Among them, 237 (41.6 %) died before discharge, and
5

---

Under review as a conference paper at ICLR 2023
Table 1: Training, validation and testing data.
Total
Alive
Death in the next 72 hours Samples (Patient)
Negative Positive Ratio (Patient)
Samples (Patient)
Samples (Patient)
Train
9979 (432)
7652 (388)
2327 (149)
3:1 (3:1)
Valid
642 (24)
519 (22)
123 (9)
4:1 (2:1)
Test
2712 (114)
2187 (104)
525 (31)
4:1 (3:1)
All
13333 (570)
10358 (514)
2975 (189)
3:1 (3:1)
333 (58.4%) survived. Patients are excluded if they were diagnosed with end-stage kidney disease (ESKD) before or at the time of hospital admission, are recipients of a kidney transplant, or has RRT less than 72 or greater than 2,000 hours.
Data features include 12 temporal features (systolic blood pressure, diastolic blood pressure, serum creatinine, bicarbonate, hematocrit, potassium, bilirubin, sodium, temperature, white blood cells (WBC) count, heart rate, and respiratory rate) and 11 static features including demographics and comorbidities (age, race, gender, admission weight, body mass index (BMI), Charlson comorbidity score, diabetes, hypertension, cardiovascular disease, Chronic Kidney Disease, and Sepsis). All out- liers (> 97.5% or 2.5%) were excluded and missing values are imputed with the last observation carried forward (LOCF) method.
To continuously predict mortality risks, we generate 30 samples from each patient's EHR data with random start and end times as long as the duration is greater than 10 time steps. The class label of sample whether the patient died (positive) or survived (negative) in the next 72 hours. From 570 AKI-D patients, 13,333 EHR samples are extracted, including 2.975 positive and 10,358 negative samples. As shown in Table all the EHR samples are split into train (75%), validation (5%), and test data (20%) patient-wise, which ensures that samples from the same patients only appear in one of the three sets. Eighteen subpopulations were considered in this study based on nine domains according to patient demographies (i.e., age. gender. race) and commodities (i.e., Charlson score, diabetes, hypertension, cardiovascular disease, chronic kidney disease, and sepsis).
3.2 BASELINE METHOD AND FAIRNESS PERFORMANCE METRICS
We compared MTATE with two widely used and accessible sequence DL methods, LSTM and Transformer, one well-known EHR cHe-specific representations method, RETAIN, and one pioneer domain-adaptation method, D For Transformer, the encoder part of the original Transformer is used. For D we only use the gradient reversal layer from the original DANN to get domain- invariant representation with all other structures the same as MTATE. For all models, the input data are the EHR temporal features, and static features are concatenated with latent representation be- fore the prediction layer, as described in Section 2.6 We evaluate the performance of all models using traditional performance metrics: Area under the ROC Curve (ROCAUC), Accuracy(ACO), Area under the Precision-Resdd Parity Difference (DPD), Equality of Opportunity Difference (EOD) and Equalized Odds Differ- ence (EQOD) (Feldman et al.) 2015)) |Hardt et al. 2016: Afrose et al., (022) (see fairness metric equations [5 [7] in Appendix). We compare all models on both imbalanced and balanced sets. The positive samples (died) and negative samples (survived) are 1:4 and 1:1, respectively.
Curve (PRAUC) as well as three faimess metrics Demographic
4 RESULTS AND DISCUSSION
4.1 COMPARISON WITH BASELINE METHODS
The overall performance of rolling mortality prediction in the next 72 hours for all test data with an imbalanced positive to negative ratio are shown in Table 2 We also show the performance on the balanced data in Table A1 Regarding the imbalanced test data, Table 2 shows that MTATE outperforms all the compared baseline methods in almost all metries. LSTM is the most competitive method since has the same highest ROCAUC as MTATE and the highest PRAUC. Nevertheless, the fairness scores of LSTM are not as good MTATE. RETAIN, and Transformer have a moderate

---

Under review as a conference paper at ICLR 2023
Table 2: Performance comparison mortality prediction in the next 72 hours imbalanced test data (pos:neg=1:4). DPD, EOD, and EQOD are the lower the better.
Method
ROCAUC
ACC
PRAUC
DPD
EOD
EQOD
Transformer
0.71(0.08)
0.69(0.09)
0.39(0.17)
0.18(0.10)
0.08(0.08)
0.10(0.05)
LSTM
0.72(0.11)
0.77(0.07)
0.55(0.20)
0.12(0.08)
0.10(0.07)
0.07(0.04)
RETAIN
0.69(0.12)
0.78(0.07)
0.45(0.20)
0.13(0.12)
0.10(0.07)
0.07(0.06)
DANN
0.60(0.12)
0,64(0.14)
0.31(0.16)
0.23(0.18)
0,08(0.06)
0.12(0.08)
MTATE (Ours)
0.72(0.09)
0.81().07)
0.49(0.18)
0.09(0.07)
0.07(0.06)
0.05(0.03)
Transforer
LSTM
RETAIN
DANN
MTATE(Ours)
Equalizoon Oods Dimerence (EGOD)
0.25
0.20
0.10
0.05
Age
Gender
Race
CCI
DB
HT
CVD
CKD
Sepsis
AVG
Subpopulation Domain
Figure 3: Equalized odds difference for every subpopulation domain. Y-axis represents the EQOD score (the lower the better) and X-axis represents the subpopulation domains, where each domain consists of two subpopulations (e.g.. Young (< 65 y/o) vs. Old in Age and Sepsis vs. Non-Sepsis). Here, CCI, DB, CVD, and CKD stand for Charlson comorbidity score, diabetes, hypertension, car- diovascular disease, and chronic kidney disease, respectively.
performance. The performance of all the compared algorithms on the balanced test data shows that MTATE has the best ROCAUC, DPD, EOD, and EQOD, and second- beat accuracy (see Table Al..
We compare MTATE with all baseline methods within each subpopulation domain. Figure shows that MTATE has the best (lowest) averaged EQOD score. We also compare MTATE with all base- lines regarding the difference of PRAUC within each subpopulation domain (e.g., the difference of PRAUC between female and male). The difference in PRAUC for each domain and the averaged score across all domains Figure A2 shows that MTATE has the lowest percentage difference in PRAUC between subpopulations in Age, Gender, and Hypertension domains.
4.2 ABLATION STUDY
We conduct an ablation study to test how each component MTATE performs by removing them from MTATE. Table 3ishows that MTATE has the best performance for almost all metrics except for ROCAUC, which is the second best to MTATE without masking. However, the accuracy of MTATE without masking is 18% lower than MTATE. In addition, RW-ATT is the most effective component since the performance drops the most in the two ablations that are without RW-ATT.
A primary goal of MTATE is to learn fair representations that can be used by a wide spectrum of downstream predictive models. To this end, we test whether the representation learned from MTATE can be used by traditional machine learning methods. The last three lines in Table 3 shows that all three traditional methods, XGboost, SVM, and Random Forest, have achieved similar performance MTATE and are better than some of the compared deep learning methods. This comparison confirms that MTATE can serve as a pre-trained EHR data representation generator, and the learned representations can be used by downstream prediction tasks implemented with different classifiers.
7

---

Under review as a conference paper at ICLR 2023
Table 3: Performance comparison of MTATE and its ablation components for the mortality predic- tion in the next 72 hours (pos:ee = 1:4). w/o. DC: remove all domain classifiers. w/o. RW-ATT: remove representation-wise attention. w/o. DC & RW-ATT: remove both domain classifier and representation-wiss wise attention. w/o. Masking: remove masking layers in FR-Attention module.
w/o. L. remove contrastive loss.
Method
ROCAUC
ACC
PRAUC
DPD
EOD
EQOD
w/o.
0.70(0.09)
0.70(0.07)
0.47(0.18)
0.15(0.12) 12)
0.09(0.07)
0.09(0.06)
w/o. RW-ATT
0.64(0.13)
0.72(0.10)
0.37(0.21)
0.23(0.20)
0.12(0..0)
0.13(0.09)
w/o. DC & RW-ATT
0.63(0.13)
0.70(0.11)
0.37(0.21)
0.25(0.20)
0.11(0.10)
0.13(0.09)
w/o. Masking
0.73(0.09)
0.63(0.06)
0.48(0.18)
0.15(0.12)
0.09(0.08)
0.09(0.06)
w/o. Le
0.71(0.09)
0.77(0.07)
0.44(0.17)
0.11(0.11)
0.08(0.06)
0.07(0.05)
MTATE
0.72(0.09)
0.81(0.07)
0.49(0.18)
0.09(0.07)
0.07(0.06)
0.05(0.03)
XGBoost MTATE
0.69(0.09)
0.70(0.07)
0.43(0.17)
0.09(0.08)
0.07(0.07)
0.07(0.03)
SVM w/ MTATE
0.71(0.10)
(0.6))
0.48(0.1) 19)
0.10(0.08)
0.07(0.06)
0.06(0.04)
RF w/ MTATE
0.72(0.09)
0.80(0.06)
0.49(0.17)
0.07(0.07)
0.06(0.06)
0.05(0.03)
0.45
S 0.35
0.10
0.35
0.60
0.10
0.25
0.40
0.10
0.15 La (Race)
0.20
Lg (Age)
La (CVD)
Figure 4: Relationship between outcome loss, domain loss and representation-wiss attention. Y- axis represents the outcome loss, X-axis represents the domain loss. The colored dots represent representation-wise attention, and the darker color represents higher attention.
4.3 EFFECTIVENESS ASSESSMENT OF RW-ATTENTION
Since RW-Attention is the most effective component in MTATE, study its relationships with the outcome prediction loss Lp and domain loss Lg, as shown in Figure 4 Each dot in the figure presents the averaged value of all samples from the same subpopulation. The figure shows three example domains in two facets. First, the correlation between the outcome prediction loss and the domain loss could be negatively correlated (Age), positive (CVD), or mixed (Race). In the negative correlation scenario, the higher the domain loss, the lower the outcome prediction loss, which sug- gests the RW-Attention putting more weight on the representations with larger domain loss (i.e., Nomain-invariaa representations). In the positively correlated scenario, outcome prediction loss de- creases with the decrease of domain loss. It suggests the RW-Attention putting more weight on the representations with smaller domain loss (domain-specific representations). Second, attention weights, as indicated by the color of the dots in the figure, demonstrate the relationship between RW- Attention and the outcome prediction loss. The darker colored dots (greater attention) are almost always related to the lower outcome prediction loss, whether the outcome prediction loss and the do- main loss positively negatively correlated. This indicates that the RW-Attention module can weigh both domain specific or domain-invariant representations toward lower outcome prediction loss. Similar patterns in all other domains are in Figure A4
The
8

---

Under review as a conference paper at ICLR 2023
5 CONCLUSION
In this work, we presented MTATE, an attention-based encoder for EHR data. MTATE uses three different attention mechanism (time-relevance, feature-relevance,, and representation-wee wiise) to learn unbiased data representations. Experiments on real-world healthcare data demonstrated that MTATE outperforms the compared baseline methods on continuous mortality risk prediction for critically ill AKI-D patients regarding fairness.
REFERENCES
Robert Adragna, Elliot Creager, David Madras, and Richard Zemel. Fairness and robustness in invariant learning: A case study in toxicity classification. arXiv preprint arXiv:2011.0648. 2020.
Sharmin Afrose, Wenjia Song, Charles B Nemeroff, Chang Lu, and Danfeng Daphne Yao. Subpopulation-speccfiic machine learning prognosis for underrepresented patients with double pri- oritized bias correction. Communications medicine, 2(1):1-14, 2022.
Henrique Aguiar, Mauro Santos, Peter Watkinson, and Tingting Zhu. Learning of cluster-based fea- ture importance for electronic health record time-series. In International Conference on Machine Learning. pp. 161-179. PMLR, 2022.
Muhammad Aurangzeb Ahmad, Arpit Patel, Carly Eckent, Vikas Kumar, and Ankur Teredesai. Fair- ness in machine learning for healthcare. In Proceedings of the 26th acm sigkdd international conference on knowledge discovery data mining, pp. 3529-3530, 2020.
Hsin Hiiin Chang, Jung- Msien Chiang, Chi-Shiang Wang, Ping-Fand Fing-ang Chiu, Khaled Abdel-Kader, Huiwen Chen, Edward D Siew, Jonathan Yabes, Raghavan Murugan, Gilles Clermont, et al. Pre- dicting mortality using machine learning algorithms in patients who require renal replacement therapy in the critical care unit. Journal of Clinical Medicine, 1 1(18):5289, 2022.
Alexandra Chouldechova and Aaron Roth. The frontiers fairness in machine learning. arXiv preprint urXiv:1810.088// 2018.
Junyoung Chung, Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron C Courville, and Yoshua Ben- gio. A recurrent latent variable model for sequential data. Advances in neural information pro- cessing systems, 28, 2015.
Michael Feldman, Sorelle A Friedler, John Moeller, Carlos Scheidegger, and Suresh Venkatasubra- manian. Certifying and removing disparate impact. In proceedings of the 21th ACM SIGKDD international conference on knowledge discovery and data mining, pp. 259-268, 2015.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversariall training of neural net- works. The journal of machine learning research, 17(1):209-22)) 2030, 2016.
Sara Gerke, Timo Minssen, and Glenn Cohen. Ethical and legal challenges of artificial intelligence- driven healthcare. In Artificial intelligence healthcare, pp. 295-336. Elsevier, 2020.
Milena A Gianfrancesco, Suzanne Tamang, Jinoos Yazdany, and Gabriela Schmajuk. Potential bi- ases in machine learning algorithms using electronic health record data. JAMA internal medicine, 178(1 1):1544-1547 2018.
Moritz Hardt, Eric Price, and Nati Srebro. Equality of opportunity in supervised learning. Advances in neural information processing systems, 29, 2016.
Farzanch Khoshnevisan and Min Chi. Unifying domain adaptation and domain generalization for robust prediction across minority racial groups. In Joint European Conference on Machine Learn- ing and Knowledge Discovery in Databases, pp. 521-537. Springer, 2021.
Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishman. Supervised contrastive learning. Advances in Neural Information Processing Systems, 33: 18118866118677 2020.

---

Under review as a conference paper at ICLR 2023
Changhee Lee and Mihaela Van Der Schaar. Temporal phenotyping using deep predictive cluster- ing of disease progression. In International Conference on Machine Learning, pp. 5767-5777. PMLR, 2020.
Hyeon-Ju Lee and Youn-Jung Son. Factors associated with in-hospital mortality after continuous renal replacement therapy for critically ill patients: a systematic review and meta-analysis. Inter- national journal of environmental research and public health. 17(23):8781, 2020.
Ninareh Mehrabi, Fred Morstatter, Nripsuta Saxena, Kristina Lerman, and Aram Galstyan. A survey on bias and fairness in machine learning, ACM Comput. Surv. 54(6), jul 2021. ISSN 0360-0300. doi: 10.1145/345767, URL httptttttpooiiii.. /ttp///....rrr111.11/3333335
Javier A Neyra and Girish N Nadkarni. Continuous kidney replacement therapy the future: In- novations in information technology, data analytics, and quality assurance systems. Advances in Chronic Kidney Disease, 28(1):13-19, 2021.
Trishan Panch, Heather Mattie, and Rifat Atun. Artificial intelligence and algorithmic bias: impli- cations for health systems. Journal of global health, 9(2), 2019.
Ravi Parikh, Stephanie Teeple, and Amol S Navathe. Addressing bias in artificial intelligence in health care. Jama, 322(24):2377-2375 2019.
Sanjay Purushotham, Wilka Carvalho, Tanachat Nilanon, and Yan Liu. Variational recurrent adver- sarial deep domain adaptation. In International Conference on Learning Representations, 2017.
URL httttss:s//ppeeeeeettettt
Forum? iidrk htt.:.net/ooorm.eet/
9
Farah Shamout, Tingting Zhu, and David A Clifton. Machine learning for clinical outcome predic- tion. IEEE reviews in Biomedical Engineering. (4:116-126, 2020.
Suvi T Vaara, Pavan K Bhatraju, Natalja L Stanski, Blaithin A McMahon, Kathleen Liu, Michael Joannidis, and Sean M Bagshaw. Subphenotypes acute kidney injury: narrative review. Critical Care, 26(1):1-10, 2022.
Cao Xiao, Edward Choi, and Jimeng Sun. Opportunities and challenges in developing deep learning models using electronic health records data: a systematic review. Journal of the American Medical Informatics Association, 25(10):1419-1428 2018.
10

---

Demographic parity suggests that a predictor is unbiased if the prediction is independent of the protected attribute (e.g., Age, Gender, and etc.). We denote protected attribute as E a, b, and A only take two groups b (e.g., Young vs Old for Age) for simplicity. Thus, the Demographic parity Under review as a conference paper at ICLR 2023
A APPENDIX
A.1 FR-ATTENTION
The FR-Attention sub-module for each head is computed
Q' K V MUq. MUK. K. MUv
(10)
A FR softmar( Q'KT
(11)
V
M V
(12)
M' Concat (M .... Mtt M )Uo
(13)
Z' max(0, (M')TW
W2 + b2
(14)
Similar to TR-Attention module, we assume all projection matrix Uq. UK, Uv have the same di-
mension d'.. Thus, Uq. UK. Uv E R.x. and Q', K', V E Rixx, The feature-relevance attention A FR & RM The output of each head is M E R Nexi,, Similar to the TR-Attention module, all outputs from all heads are concatenated to form M' E R RNx,,, and linear transforma- tion are applied with the projection matrix Uo E IR xM where N₁, represents the number of head.
A.
FR-Attention
B.
Multi-Head Attention
Add & Norm
Linear
FFN
Concat
Transpose
M'
Add Norm
Scaled Dot-Product
Attention
Multi-Head Attention
Mask
Linear
Linear
Linear
z
Figure A1: A. The structure of FR-Attention module in MTATE. B. The multi-head attention module from the original Transformer used in MTATE.
A.2 FAIRNESS METRICS
difference (DPD) is the difference between the two group a and b The formula for Demographic parity difference (DPD) is shown in below:
DPD P 1|A a) P(ğ 1/A b)
(15)
Equality of opportunity suggests that predictor is unbiased if the true-positive rate between two groups are equal. Similarly, the Equality of opportunity difference (EOD) is the difference between the two group a and b. The formula for EOD is shown in below:
EOD P(j 1ly 1, = a) P(ù 1/y = 1, b)
(16)
11

---

Under review as a conference paper at ICLR 2023
Equalized odds suggests that predictor is unbiased if both the true-positive rate (TPR) and false- positive rate (FPR) between two groups are equal. We computes the Equalized odds Difference (EQOD) the average of the difference in both TPR and FPR. The formula for EQOD is shown in below:
EQOD = (TPRp + FPRp)/2 =
(17)
TPRD P(ù 1
11y=11y 1,, b) P(ÿ 11y - 0, =b)
(18)
FPRD PÚù =
(19)
A.3 PERFORMANCE COMPARISONS
A.3.1 OVERALL PERFORMANCE
The following table shows the performance comparison between MTATE and baseline methods for the balanced test data.
Table A1: Balanced performance of MTATE and compared algorithms for the mortality prediction in the next 72 hours (pos:neg 1:1).
Method ROCAUC
ACC
PRAUC
DPD
EOD
EQOD
Transformer
0.70(0.09)
0.67(0.09)
0.69(0.11)
0.17(0.13)
0.15(0.12)
0.10(0.06)
LSTM
0.71(0.11)
0.67(0.11)
0.76(0.12)
0.19(0.12)
0.20(0.12)
0.12(0.05)
RETAIN
0.68(0.12)
0.67(0.11)
0.72(0.13)
0.22(0.13)
0.21(0.12)
0.12(0.06)
DANN
1.58)
0.54(0.14)
0.59(0.12)
0.21(0.19)
0.16(0.13)
0.12(0.08)
MTATE (Ours)
0.73(0.09)
0.65(0.11)
0.75(0.11)
0.15(0.10)
0.14(0.11)
0.08(0.05)
A.3.2 FAIRNESS PERFORMANCE WITHIN EACH DOMAIN
1.2
MTATEICMMS)
1.0
0.8
0.4
Age
Gender
Race
OCI
D8
HT
CVD
CKD
Sepeis
AVG
Subpoputation Domain
Figure A2: The comparison between MTATE with baseline methods for the percentage difference score PRAUC for each domain. Y-axis represents the percentage difference. X-axis represents the subpopulation domain, each domain consists of two subpopulations (c.g., Young (< 65 y/o) Old in Age domain, Sepsis vs. Non-Sepsis in Sepsis Domain ). CCI stands for charlson comorbidity score, DB stands for diabetes, HT stands for hypertension, CVD stands for cardiovascular disease, CKD stands for chronic kidney disease.
12

---

Under review as a conference paper at ICLR 2023
A.4 STUDY OF MASKING RATE
We analysed the effect on the performance by different masking ratio. Figure A3 shows the per- formance metrics of ROCAUC, PRAUC and accuracy with respect the masking ratio from 0 to 0.9. Both ROCAUC and PRAUC have a similar trend, both starting with a relative high score and gradually decrease with some punctuation, and both reach lowest scores masking rate 0.8 and 0.9. In contrast, the accuracy has opposite trend, where it starts with the lowest score, and almost monotonically increasing. The masking rate highly dependant on the input data, thus the chosen of masking rate is not universal. For our experimental data, one of the best masking rate is 0.4, where it has the highest accuracy and PRAUC and the third best ROCAUC.
0.8
0.8
0.9
0.81
RODAUC
0.73
0.710.71
0.70
0.71
0.72
PRAUIC
0.5
0.46
0.49
0.49
0.8
0.75
0.73
700
0.7
0.7
0.67
9/41
0.40
0.88
66630.00 63
0.8
0.8
0.3
0.5
0.5
Mosking Hate
Mosking Kate
Masking Hato
Figure A3: Performance Score of ROCAUC, PRAUC and Accuracy with different masking rate
A.5
RELATIONSHIP BETWEEN OUTCOME LOSS, DOMAIN LOSS AND REPRESENTATION- WISE ATTENTION
044
-3
0.3
0.10
0.35
0.80
0.2
0.4
0.8
0.10
0.15
0.20
0.10
0.36
0.60
0.1
0.3
0.5
(Aga)
La (Candar)
(Race)
(0D)
(g(RR)
-3
0.2
0.40
0.10
0.50
La (HT)
4a (CVD)
La (CKD)
La (Scpsis)
Figure A4: Relationship between outcome loss, domain loss and representation-wiss attention in all domains. Y-xxis represents the outcome loss, x-axis represents the domain loss. The colored dots represent the representation-wiin attention, and darker color represents higher attention.
13


Reviews:
[{"review_id": "r8uulO_YumE", "reviewer": "Reviewer_M4Hz", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "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": "", "summary_of_the_paper": "The authors present a transformer-encoder framework for multipleThe authors present a encoder/transformer method for learning representations of EHR data,  for multiple prediction tasks. They demonstrate that their method out-performs other DL methods on one prediction task using MIMIC-III data.\n", "summary_of_the_review": "This is a well-written paper on an interesting and important topic. The experiments are somewhat compelling andt he methods appear to be sound. The main issue is that the authors don't disucss or compare with other methods from the literature. There has been a lot of work in the encoder/transformer space for EHR data, and the authors don't engage with this literature. Since they are presenting their method as a general approach for EHR data (and not a method specific to AKI), they should make sure to engage & compare with exisitng work.\n", "strength_and_weaknesses": "\nStrengths:\n- This paper is well written\n- it focuses on an important problem (they also motivate the problem very well)\n- the methods seem sound, though I am not an expert\n- the experiments use relevant data, and the results demonstrate that their method works well\n- their experiments use public data\n\nWeaknesses\n- There has been a lot of work on this topic, and the authors don't cite or compare with existing methods. Here is a short list of papers that look relevant to this topic, but were not mentioned or compared with in the paper:\n-- https://doi.org/10.1109/ICDM50108.2020.00050\n-- https://doi.org/10.1145/3394486.3403129\n-- https://doi.org/10.1109/ICDM51629.2021.00060\n-- here is a relevant review: https://doi.org/10.1016/j.jbi.2020.103671\n- there is no comparison against state-of-the-art risk prediction models in medicine, or any non-DL models at all. there are several such models used in medicine (which don't use DL), which would be good baselines to compare against here. \n- the authors don't provide any code\n\n", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "nkBYj0zgnt", "reviewer": "Reviewer_icM1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "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": "", "summary_of_the_paper": "The authors propose a method to learn representation that is invariant to sensitive groups (or domains as claimed by the authors). This is achieved through a network structure that has a transformer to generate time-wise attention, domain-specific feature-wise attention (to extract features that are specific to each domain/subgroup), and finally all these representations are concatenated and passed to a network to generate a universal representation that can be used for downstream tasks.", "summary_of_the_review": "Read the strength and weakness section. ", "strength_and_weaknesses": "Strength: \n- The problem is very important especially in health care domain. \n\nWeakness: \n- The paper is not easy to follow \n- I think using the word \"domain\" might by non-intuitive in this context. It is usually referred to as sensitive or protected groups. \n- The method seems to be adhoc without a principle way to motivate the structure of the network. \n- It is claimed in the paper that \"To the best of our knowledge, MTATE is the first model that automatically trains and determines optimal and fair data representations.\" There are many papers about learning fair representation. The one that needs to be highlighted here is \"Creager, Elliot, et al. \"Flexibly fair representation learning by disentanglement.\" International conference on machine learning. PMLR, 2019.\" The method in that paper learns representation that is invariant to the sensitive groups, which is the main objective of the current paper. \n- Missing important baselines to compare to. ", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "UwTet9ACLs", "reviewer": "Reviewer_UWzS", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "", "summary_of_the_paper": "In this paper, the authors propose an attention-based encoder model to learn data representations that are fair for prespecified subpopulations. The proposed method was tested on real data to learn EHR representations and predict 72-hour mortality after dialysis.", "summary_of_the_review": "The method in the paper is proposed to balance domain-specific and domain-invariant representations with the aim of learning fair and unbiased representation. The experiment has shown its advantage empirically. Overall, I think the proposed method could be useful. But more details can be provided in the paper.", "strength_and_weaknesses": "The experiment was well conducted. The proposed model has shown advantages in both prediction performance and fairness across prespecified subgroups. \n\nMore explanation and intuitions can be provided for the proposed method, for example, the rationale for introducing each model part and their connection with existing methods. Also, the reason why the model can achieve desired fairness can be better described.\n\n", "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": "JzAiM8CHJrH", "reviewer": "Reviewer_H5nZ", "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": "", "summary_of_the_paper": "This paper addresses the fairness problem in a healthcare setting involving outcome prediction from electronic health records (EHR). To this end, the authors propose the Masked Triple Attention Transformer Encoder (MTATE), an adaptive multi-task learning model capable of automatically learning and selecting optimal yet fair data representations without explicit domain adaptation or domain-specific bias correction. The experimental results suggest that MTATE learns representations of diverse subpopulations that lead to improvements in predictive performance and fairness on the task of mortality risk prediction for critically ill Acute Kidney Injury requiring Dialysis (AKI-D) patients.", "summary_of_the_review": "Overall, this paper proposes an adaptive multi-task learning model capable of learning optimal and fair data representations without explicit domain adaptation or domain-specific bias correction. Addressing fairness in a healthcare setting involving patient outcome prediction is of considerable importance. Nevertheless, the work lacks methodological novelty and is rather incremental in that respect considering that it combines already established attention mechanisms. That being said, I recommend that this work is not considered for acceptance at ICLR due to reasons including, but not limited to, the following:\n\n(1) the degree to which this work is application-oriented rather than a work that makes novel methodological contributions in representation learning;\n\n(2) the absence of a comparison of the proposed model to more recently published models for representation learning from EHR data;\n\n(3) unjustified data preprocessing decisions (outlined in the list of weaknesses);\n\n(4) experiments conducted solely on proprietary data despite the availability of public EHR datasets that also contain records of AKI-D \npatient among other conditions.\n\nDespite the above reasons, I am still looking forward to the authors’ response and I would be willing to adjust my score in case I have misunderstood or misinterpreted certain aspects of the work.", "strength_and_weaknesses": "Strengths:\n\n* The problem addressed in this paper is of considerable importance as approaches that account for patients’ backgrounds and subpopulations are quite relevant to personalized medicine and reliable prediction of patient outcomes.\n\n* The proposed MTATE model is capable of learning optimal yet fair representations of EHR sequences, which is of great importance to outcome prediction for patients with different demographic characteristics and clinical histories.\n\n* MTATE incorporates time-wise, feature-wise and representation-wise attention mechanisms. This allows for (1) capturing associations among timesteps (which are quite informative in the case of EHR data), associations among features, and (2) summarizing feature-attended representations for downstream patient outcome prediction.\n\n* While it is unclear whether domain-invariant or domain-specific data representations should be used for prediction tasks concerning EHR data, the authors claim that MTATE learns EHR representations that can be domain-specific, domain-invariant, or even a mixture of the two.\n\n* Experiments conducted on a real-world EHR dataset demonstrate that MTATE outperforms several outcome prediction baselines on mortality risk prediction for critically ill AKI-D patients, yielding improvements in both classification metrics as well as fairness metrics, thus mitigating bias towards different subpopulations in the data.\n\n* A few traditional classifiers (XGboost, SVM and RF) have achieved performance comparable to that of MTATE and outperformed some of the compared baseline methods; which suggests that MTATE’s representations can also serve as pre-trained EHR data representations for various classifiers applied to downstream prediction tasks.\n\n* The authors conducted an insightful analysis of the relationship between the patient outcome loss, the domain loss and the representation-wise attention.\n\n-------------------------------------------------------------------------------------------------\n\nWeaknesses:\n\n* The proposed MTATE model appears to resemble a multi-task attention classifier with an output layer that combines the representations learned for all ‘tasks’ (i.e. domains) into a single outcome. Due to this resemblance, I believe that a comparison of MTATE to multi-task classification models should have been carried out. I would encourage the authors to consider such a comparison or provide their reasoning as to why multi-task classification models should not be considered among the baselines.\n\n* In Section 2.3, the authors state the following: “Since features for the different domains are not equally important, we randomly masked out some number of latent features along all time steps differently for each FR-Attention module. In return, the masking procedure forces each sub-module to learn different feature focuses for each domain to generate unique domain representations.” However, no experiment has been conducted to inspect the impact of the random masking on the learned feature representations and thus on the overall patient outcome prediction performance. I am aware of the ablation study (Table 3) conducted to assess the performance with and without masking, nevertheless, such a study does not provide an insight into the effect of different random masking trials. I would suggest that the authors elaborate on this aspect of the random marking in their response.\n\n* One of my main concerns is that the resulting representations learned for each domain can be domain-specific or domain-invariant, depending on the domain classification loss. Thus, target domains that are “easier to classify” will tend to have smaller loss values. On the other hand, other target domains might be difficult-to-classify yet contain more informative patterns, in which case such domains will have higher loss values and may need many iterations for their losses to be substantially minimized. This case does not appear to be accounted for in the model design. I would encourage the authors to clarify the potential behavior of MTATE in the aforementioned scenario.\n\n* In Section 2.6, the authors mention that “the final representation C of EHR is concatenated with all static features, such as demographics and comorbidity, followed by a linear layer with a sigmoid function for the outcome binary prediction”. I assume that the demographics and comorbidity features are one-hot encoded, in which case it would be desirable that the range of the learned representations’ values is around [0, 1]. Is a certain normalization technique applied to the final representations to scale them within the same (or similar) range as that of the static features? If this is not the case, I would like to ask the authors to provide their thoughts on this point in their response.\n\n* The authors have generated 30 samples from each patient’s EHR history with a random start and end times. Nevertheless, I am wondering whether there is a specific reason behind this data preprocessing choice. Why not consider the entire EHR sequence for each patient? Were the sequence samples generated just for data augmentation purposes, i.e. so that representations can be learned from a larger dataset created by generating multiple samples from a single EHR sequence?\n\n* In Section 3.2, the authors mention that all models have been compared on both imbalanced and balanced sets having positive/negative ratios of 1:4 and 1:1, respectively. However, I am not certain if 1:4 can be considered imbalanced (of course, the ratio itself is not a balanced ratio, but the classification problem cannot be considered a typical case of an imbalanced classification problem, or at least not a highly imbalanced classification problem).\n\n* In addition to the considered baselines, I believe that MTATE should have been compared to several recently published state-of-the-art representation learning methods for EHR data including, but not limited to, the methods proposed in:\n\n    Darabi, S., Kachuee, M., Fazeli, S., & Sarrafzadeh, M. (2020). Taper: Time-aware patient ehr representation. IEEE journal of biomedical and health informatics, 24(11), 3268-3275.\n\n    Bang, S. J., Wang, Y., & Yang, Y. (2020). Phased-lstm based predictive model for longitudinal ehr data with missing values.\n\n    Bai, T., Zhang, S., Egleston, B. L., & Vucetic, S. (2018, July). Interpretable representation learning for healthcare via capturing disease progression through time. In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining (pp. 43-51).\n\n-------------------------------------------------------------------------------------------------\n\nMinor weaknesses:\nThere are also certain grammatical and typographical errors and remarks that require attention. Some of them are summarized as follows:\n- First paragraph on page 4: The sentence “Thus, the attention weights computed from the query and key represent how much focus each time step is associated with other time steps.” is rather unclear and should be revised accordingly.\n- Page 4: In the last sentence of Section 2.4, “generate” should be replaced with “generating” and a word seems to be missing between “for” and “in”.\n- Page 5: “C” should be bolded in the paragraph after Eq. (9).\n- Last paragraph on page 6: In the sentence “LSTM is the most competitive method since it has the same highest accuracy”, I assume that the authors meant “same highest ROCAUC” instead of “same highest accuracy”.\n- Page 7, last paragraph of Section 4.1: I believe that “Figure 4.1” should be replaced with “Figure 3”. Later in the same paragraph, “Figure A.3.2” should be replaced with “Figure A4”.\n- In Section 4.3: replace “negative correlated” with “negatively correlated”.\n", "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": {}}]

Author rebuttal:

Decision: Reject