CURE: A Pre-training Framework on Large-scale Patient Data for Treatment Effect Estimation
Authors: Ruoqi Liu
Pin-Yu Chen
Ping Zhang
Venue: Submitted to ICLR 2023 (2023)

Abstract: Treatment effect estimation (TEE) refers to the estimation of causal effects, and it aims to compare the difference among treatment strategies on important outcomes. Current machine learning based methods are mainly trained on labeled data with specific treatments or outcomes of interest, which can be sub-optimal if the labeled data are limited. In this paper, we propose a novel transformer-based pre-training and fine-tuning framework called CURE for TEE from observational data. CURE is pre-trained on large-scale unlabeled patient data to learn representative contextual patient representations, and then fine-tuned on labeled patient data for TEE. We design a new sequence encoding for longitudinal (or structured) patient data and we incorporate structure and time into patient embeddings. Evaluated on 4 downstream TEE tasks, CURE outperforms the state-of-the-art methods in terms of an average of 3.8\% and 6.9\% absolute improvement in Area under the ROC Curve (AUC) and Area under the Precision-Recall Curve (AUPR), and 15.7\% absolute improvement in Influence function-based Precision of Estimating Heterogeneous Effects (IF-PEHE). We further demonstrate the data scalability of CURE and verify the results with corresponding randomized clinical trials. Our proposed method provides a new machine learning paradigm for TEE based on observational data. 

Under review as a conference paper at ICLR 2023
CURE: A PRE-TRAINING FRAMEWORK ON LARGE-SCALE PATIENT DATA FOR TREATMENT EFFECT ESTIMATION
Anonymous authors Paper under double-bling blind review
ABSTRACT
Treatment effect estimation (TEE) refers to the estimation causal effects, and it aims to compare the difference among treatment strategies on important out- comes. Current machine learning based methods are mainly trained on labeled data with specific treatments or outcomes of interest, which can be sub-optimal if the labeled data are limited. In this paper, we propose a novel transformer -based pre-training and fine-tuning framework called CURE for TEE from observational data. CURE is pre-trained on large-scale unlabeled patient data to learn repre- sentative contextual patient representations, and then fine-tuned labeled patient data for TEE. We design a new sequence encoding for longitudinal (or structured) patient data and we incorporate structure and time into patient embeddings. Eval- uated on 4 downstream TEE tasks, CURE outperforms the toftc--ff-h-aattt meth- ods in terms of an average of 3.8% and 6.9% absolute improvement in Area under the ROC Curve (AUC) and Area under the Precision-Recall Curve (AUPR), and 15.7% absolute improvement in Influence function-baan based Precision of Estimating Heterogencous Effects (IF-PHHE). We further demonstrate the data scalability of CURE and verify the results with corresponding randomized clinical trials. Our proposed method provides machine learning paradigm for TEE based on observational data.
INTRODUCTION
Treatment effect estimation (TEE) is to evaluate the causal effects of treatment strategies on some important outcomes, which is a crucial problem in many areas such as healthcare (Glass et al.) 2013). education (Dehejia & Wahba 1999) and economics (Imbens 2004). Randomized clinical trials (RCTs) are the de-facto gold standard for identifying causal effects through randomizing the treatment assignment and comparing the responses in different treatment groups However, con- ducting RCT's is time-consuming, expensive and sometimes unethical. Observational data such as medical claims provide a promising opportunity for treatment effect estimation when RCTs pensive or impossible to conduct.
Recently, many works have been proposed adopt neural networks (NNs) for TEE from obser- vational data (Shalit et al.] (2017] Shi et al. 2019 Hassanpour & Greiner 2019 Curth & van der Schaar 2021b al Zhang al.J 2022b Guo al. (022). Compared to classical TEE methods such as regression trees (Chipman et al.) 2010) or random forests (Wager & Athey 2018). NN-based methods achieve better performance in handling the complex and nonlinear relationships among co- variates, treatment and outcome. However, there are still some common limitations of existing TEE methods: 1) Most model designs are task-specific or data-specific so is hard to adapt the model to more generalized setting. 2) Existing labeled dataset often has small-scale data size, whereas training neural models requires large and high-quality labeled data for capturing inherent complex relationships of the input data.
Recently, Transformer (Vaswani et al. 2017) has been widely adopted as a critical and unified building block in the pre-training and fine-tuning paradigm across data modalities. The pre- trained Transformer based models (PTMs) have become the model choice in many deep learning do- mains such natural language processing (NIP) (Devlin et al. 2018 Radford et al.) 2018 2019

---

Under review as a conference paper at ICLR 2023
Downstream Data: a
Unlabeled Patient Data
Downstream
Treatment Outcome
Z
Patient Sequence
Unlabeled Data:
Unsupervised Pre-training CMIM (0)
Fine-tuning for Treatment Effect
X
Large-scale Structured
Estimation TTEE (6' .4)
28
Patient Data
Transformer Encoder
Linear Head fa
Section 3.1
Section 3.2
Section 3.3
Figure 1: The overall pipeline of CURE. mainly consists of three parts: 1) data encoding of lon- gitudinal patient data; 2) unsupervised pre-training on unlabeled data and 3) fine-tuning on down- stream labeled data for treatment effect estimation. In TEE, labels mean the studied treatment a and outcome y of the patient sequence
Brown et al. 2020; Liu et al. 2019) and computer vision (CV) (Carion et al. 2020 Dosovitskiy et al. 2020 Parmar et al. 2018). The dominant approach is to pre-train with unsupervised or self-supervised learning and then fine-tune on a smaller task-specific dataset. Nonetheless, applying this pre-training and fine-tuning paradigm to treatment effect estimation prob- lems faces the following three major challenges: 1) encoding structured longitudinal observational patient data into sequence input: 2) lack of well-curated large-scale pre-training dataset; 3) lack of real-world downstream treatment effect estimation tasks to benchmark baselines.
large- dataset
In this paper, we propose new pre-training and fine-tuning framework for estimating causal effect of a treatment: CaUsal (Reatment Effect estimation (CURE). As shown in Fig. T the large-scale structured patient data are extracted from a real-world medical claims data (MarketScan Research Databases D. We first encode the structured data as sequential input by chronologically flattening and aligning all observed covariates. We obtain around 3M processed unlabeled patient sequences for pre-training. And the downstream datasets with labeled treatment and outcome created ac- cording to specific TEE tasks from established RCTs. Based on the retrospective study design and domain knowledge, tient samples. The task is to evaluate the comparative effectiveness of two treatment effects in reducing the risk of stroke for patients with coronary artery disease (CAD). Second, pre-train Transformer-based model on the unlabeled data with an unsupervised learning objective to generate contextualized patient representations. To accommodate the issues of complex hierarchical structure (i.e.. the patient record contains multiple visits and each visit contains multiple types of medications or diagnoses) and irregularity of the observational patient data, we propose comprehensive embed- ding method to incorporate the structure and time information. Finally, we fine-tune the pre- pre-trained model on various downstream TEE tasks.
obtain 4 downstream tasks and each of them containing 10K-20K pa-
We are the first study to demonstrate the success of adopting the pre- teeiiraining and fine-tuning tuning frame- work to representation learning of patient data for TEE, together with necessary but minimal changes on the transformer architecture design, and real-world case studies on randomized clinical trials. We summarize our main contributions as follows.
We propose CURE, novel transformer-basee pre-training and fine-tuning framework for TEE. We present a new patient data encoding method to encode structured observational patient data and incorporate covariate type and time into patient embeddings.
We obtain and preprocess large-scale patient data from real-world medical claims data as pre- training resource. We derive downstream TEE tasks according to study designs and domain knowledge from established RCTs for model evaluation.
We conduct thorough experiments and show that CURE yields superior performance on all down- stream tasks compared to srate-of-the-ee TEE methods. We achieve, on average, 3.8% and 6.9% absolute improvement in AUC and AUPR respective for outcome prediction, and 15.7% abso- lute improvement in IF-PEHE for TEE over the best baseline among 4 tasks. We also verify the estimated treatment effects with the conclusion of corresponding RCTs.
We further explore the effectiveness of CURE in several ablation studies including the proposed patient embedding, the influence of pre- data size on downstream tasks, and the generaliz- ability of low-resource fine-tuning data.
2

---

Under review as a conference paper at ICLR 2023
2 BACKGROUND AND RELATED WORK
Treatment effect estimation from observational data. In this paper, we are interested in obser- vational patient data. Each patient sample consists of pre-treatment covariates ø (i.e., historical co-medication, co-morbidities and demographics) and a treatment a of interest. Following the po- tential outcome framework (Rubin) 2005), the potential outcome yo is defined as the response to treatment a out of all available treatment options. Typically, we consider the comparative treatment effects of two treatments and denote two potential outcomes as y/1 and yo for simplicity.
We aim to estimate the individual-level treatment effect (TTE) as the difference between the potential outcomes under two treatment arms as 1/1 (z)-N0(<))) We are also interested in the average treatment effect (ATE) which is the average effect among the entire population, denoted as (Byy1(z) yo (-))] In observational data, only one of the potential outcomes available and the remaining counterfac- tual outcomes are missing in nature, which makes this task more difficult than classical supervised learning. We follow the standard assumptions (Imbens & Rubin 2015) (i.e., consistency, positivity and strong ignorability). The potential outcome can be defined yo (z) E[y a, æ1,, and can be estimated from observational data. More details of assumptions are illustrated in Appendix [A]
Deep learning for treatment effect estimation. Generally, existing NN-based methods formulate the TEE as several regression tasks (i.e., regression on potential outcomes and treatment) with differ- ent levels of information shared among the nuisance estimation tasks using representation learning. TARNet (Curth & van der Schaar (2021b), for example, learns one shared representations for two potential outcomes, while SNet (Curth & van der Schaar 202 learns five different representa- tions on the combinations of treatment and potential outcomes. Recently, Transformer has been introduced an encoder block for TEE (Zhang et al., 2022b] Guo et al.) (021) and yields better performance compared to the state-of-ther the-the-art methods. Despite the promising results, the main lim- itation that the model performance can be diminished if the labeled dataset is limited. The model trained for one particular problem or data may fail to generalize to other scenarios.
Pre-train and fine-tune of Transformer. Since Transformer based on a flexible architecture with few assumptions the input data structure, it is difficult to directly train the model small-scale models (PTMs) are first pre-trained on the
data. Therefore, various pre-trained Transformer-basee large-scale unlabeled data and then fine-tuned for labeled tasks at hand. PTMs learn universal and contextualized representations, which can boost various downstream tasks, and avoid developing and training a new model from scratch. Among the existing PTMs in NLP, BERT (Devlin et al. 2018) is one of the most popular models. BERT (Devlin et al., 2018 is pre-trained on large-scale unlabeled corpus via self-supervised pre-training tasks (i.e., masked language modeling and next sentence prediction) and fine-tuned on downstream tasks with an additional linear head. Our observational patient data are close natural language text as they both contain sequential information. However, patient data have some unique characteristics that distinguish them from the text. Compared to the text, patient data contain more complex hicrarchical structure and time information. Therefore, existing BERT pre- pre-training architecture can not be directly applied to modeling patient data.
3 CURE: A PRETRAINING AND FINE-TUNING FRAMEWORK FOR TEE
In this section, we introduce our CURE framework (as shown in Fig. which includes three key steps: (1) Encoding structured patient data sequential input by aligning medications and diagnosis in each visit chronologically (Sec. 3.1.. (2) Pre-training on a large-scale unlabeled patient data by minimizing the unsupervised objective and obtaining the optimized parameters A" (Sec. (32)). and (3) Fine-tuning on small-scale labeled downstream dataset for TEE by jointly optimizing 6 and
linear head parameterized as Φ (Sec. 3.3).
3.1 ENCODING STRUCTURED PATIENT DATA
In this work, focus on longitudinal observational patient data. We first introduce the data for pre-training and fine-tuning respectively. Then we illustrate how to convert structured patient data into sequential input for the Transformer encoder.
Pre-train data structure. The pre-training is based on large-scale unlabeled patient data. Here, to distinguish from downstream data, denote the pre-train data unlabeled data (xr ~ 2), while
3

---

Under review as a conference paper at ICLR 2023
Target
Bessine pariod (++ 365 days)
Targen
anoodling
Treatment
example
Patient
oeghalasin
Medication:
digouin
montelluka lisinopel
Medication:
Medication:
treatment
group
Follow-u
age:
gender: femole
nsteoarthvitss mood disorders
Diagnosis:
chest pain paralysis
Diagnosis:
riveroxakenn
simmaszatin
docecyyiin
Potents dara
sease resaton date
Index date fsr pronooop on
v
Target) encoding
Dutcome
stan-date
a Compared)
Compared
Baseline period
Visit
Vist
treatment
Day -351
Day 33
Day -14
group
Compared anciding
Treatment
Sequenoc input
mode
cophalionn
olgoxin
paralysie
chstt pain
Jon
Index
Figure 2: Illustration of the downstream data construction with retrospective study design. Index date refers to the first prescription of the target treatment or the compared treatment, which should be no prior to the disease initia- tion date. The baseline period is no less than one year and the follow-up period as outcome observation is also one year. The treatments of interests and outcomes are obtained at the index date and follow-up period respectively.
Inpul ids
871
23-443
88
234
31
143
77
65
11
23
Illustration of encoding structured patient data into sequential input. The patient data are recorded in an hierarchical structure such that a patient contains multiple visits and each visit contains multiple medications and diagnoses. The structured data are converted into sequence by flattening all covariates in each visit and aligning them chronologically.
Figure 3:
downstream data with treatment a and outcome y labeled data (fx,, a,z,a,,, v. 2). The unlabeled patient data consist of: (1) Co-medication m, m......, , Mi E .M, where [M]] is the number of
unique medication names. (2) Co co-mobbidities d1, d2.... dIdi E D, where [D] is the number of unique diagnosis codes. (3) Demographies age encoded as categorical value and gender encoded as binary value. A patient can have multiple visits (v1....., (r), where each of visit e contains a subset of medication and diagnosis codes (w E .M D). We denote the unlabeled patient data as
x = {c, {er (=)) and all the covariates are obtained from the baseline period as shown in Figure We build medical vocabulary from all patient covariates as V (M, D,c).
Fine-tune data structure. The fine-tuning is based on small-scale labeled patient data, which are not used for pre-training. Besides the co-medication, co-morbidities and demographics, the labeled patient data contain treatment a E Mask (i.e., can be either the target treatment or com- pared treatment from task task-specific medication group Mus)) and outcome ya E {0, 1} under the observed treatment a. In Figure. 2 we show the retrospective study design of how to construct downstream data and obtain labels for treatments and outcomes. Specially, we collect patient data from two different treatment groups for comparison. For each group of patients, the covariates (i.e., co-medication, co-morbidities and demographics) are obtained from the baseline period (a.k.a., pre- treatment covariates) as potential confounders, and the outcomes obtained from the follow-up period. More illustrations of the study design can be found in Appendix [
Structured patient data to sequential input. As introduced above, the original patient data recorded naturally in a hierarchical structure. Unlike natural language text, which is inherently en- coded as sequence of words, the patient data need to be preprocessed into "sequence- format before sending to the Transformer encoder. As shown in Fig. 3 we flatten the structured patient data by chronologically going through each medication and diagnosis in each visit and aligning them in one sequence. Each medication or diagnosis is encoded as an individual token, which is comparable to text tokenization. The token ids are obtained from the medical vocabulary V.
3.2 PRE-TRAINING CURE
As shown in Fig. the pre-training consists of three modules: (1) an embedding layer to convert input patient data into embedding representations, (2) Transformer encoders to generate contextual- ized hidden representations and (3) a final project layer for pre-training objective. More formally, given the encoded patient sequence training procedure can be decomposed into the following steps:
[Z]...., Im....,, as specified in Sec. 3.1 the pre-
(I) We randomly replace 15% of input tokens with special [MASK] tokens, e.g., token Z' in the se- quence. e, RD denotes the embedding representation with embedding dimension B generated by the comprehensive embedding layer. h; E IR" denotes the contextualized representation with hidden dimension H generated by Transformer encoder fo- The masked language modeling (MLM)
Mask
FI....., [MASK] m....,
Embooking
(e.)T=. Le, (h.) MLM
LMLM(
4

---

Under review as a conference paper at ICLR 2023
Pre-training
Fine-tuning
9 (e)) (nl=))
representation
Hidden
Meausa
Hovamaaton
Hatwi
Hocesartrite
Potential Outcome Prediction
Transtorer
Pooling Layer
Type
ambedding
Taeess
Treesiuue
Freessaioo
Tyanas)
Tragrose
Nuturtan
1serm
M4324
Hawcboach
Haysape
Time
Vi Pase
Vi Preas
Vg
Passu
Vi2
Pasau
Vy Pastra
embedding
Transformer
embedking
Inpul
Wissa
Pinesater
Vironasted
Mosscart
Patient data
Demo: age 19-6
Medication
Medication wwwrkkouet
Diagnosis chout poin
Diagnosis aeboattels
Embedding Layer
encoding
Treatment
[SEP]
Demo: 45-34
Medication
Diegnossie
Raw input
Structurad Pattent Data
vallation
Invollnaain
ayyoocee
Figure 4: Illustration of pre-training and fine-tuning of CURE. The unlabeled structured patient data are first converted into sequential input, and processed for the embedding layer and encoder. During the inn-ttnnig, the treatments of interest are appended for potential outcome prediction. (Devlin et al. (2188) aims to predict the masked tokens Z'm from the established vocabulary us- ing hidden representation hmm. The pre-training loss function of MLM is denoted as (0) with optimization parameters
Comprehensive embedding layer. The pre-trained language models like BERT (Devlin et al. 2018) have achieved great success in natural language text demonstrating strong power in modeling sequential data. Though longitudinal patient data can be deemed kind of sequential data when organized by the chronological order, there exist substantial and unignorable discrepancies between the text and patient data. It is hard directly apply the existing pre-trained language model to unique patient data. Our ablation study shows that the standard embedding design adopted in NLP (i.e.. token embedding and position embedding) will be sub-optimal in our scenario (see Sec. 4.3 for more details).
Compared to the natural language text, (I) longitudinal patient data contain a more complex hierar- chical structure than the text data: a patient record contains a number of visits and each visit also contains number of different types of medical codes (i.e., medication or diagnosis). (2) The patient data are irregularly sampled (i.e., the time interval among visits is not regular) while the text data are regularly organized. As shown in Fig. 3 the visit dates are not regularly distributed along the time: the first visit happened
day 0, the second visit happened on day 14 and the third one day 33, etc. On visit 2 (day 33), the patient received two types of codes: montelukas and lisinopril as medications, and chest pain and paralysis as diagnoses.
To accommodate the above issues of complex hierarchical structure and irregularity of the obser- vational patient data, propose more comprehensive embedding layer than the original BERT (Devlin et al. 2018) embedding layer by including associated code type information and time infor- mation. For each input token, the patient embedding e, is obtained as:
Wtokenn + Leype + Uvisit + Pphysical
where Wroken is the original input token embedding. teype is the type embedding of the input token. According to our data, there are three types in total: (Demographies, Medication, Diagnosis]. The visit time embedding Vvisit is the visit time corresponding to a visit. The physical time embed- ding Pphysics: is the physical time associated with the visit. Here, the physical time is measured by month (i.e., 30-day fixed window). Both visit and physical time are organized relative to the treatment index date (i.e., the absolute distance between the visit/physical time to the index date).
As an illustration, in Fig. the input is a sequence of patient data containing the type and time information: rivaroxaban is from Medication type prescribed on visit (day 14) and chest pain from the Diagnosis type received on visit 2 (day 33). The input token embedding, time embedding and type embedding are integrated and used as the input to the Transformer encoder.
Transformer encoder and pre-training objective We use N-stacked Transformer our encod- ing backbone as it has been a widely adopted architecture. For each single Transformer encoder block, it consists of multi-head self-attention layer followed by fully-connected feed-forward layer /Vaswani et al. 2017). More details of the Transformer architecture are in Appendix B
The Transformer encoder fe takes the comprehensive embedding representations as input and gener- ates contextualized hidden representations as fo (e). Given unlabeled patient data X, the pre-training is to minimize the MLM loss of predicting the masked token with position j E using the input
5

---

Under review as a conference paper at ICLR 2023
token embedding and hidden representation:
CMLM
Ex
Σ (og(((W))o(())))
JEG
where P(wylfo(e)) is the softmax probability of the masked token over all tokens in the vocabulary.
3.3 FINE-TUNING CURE FOR TEE
Given downstream labeled data {ax o. Z, we fine-tune the model on different downstream TEE tasks. Here, we are interested in the comparative causal treatment effect of the target treatment over another compared treatment according to the downstream tasks. For each task, we plug in the task-specific imput and outputs into CURE. We add linear head fo to the hidden representations learned from the pre-training stage. We fully fine-tune all model parameters end-to-end by jointly updating f obtained from optimizing Eq. 3 and a randomized ø.
Specially, we append the original input sequence with the index treatment (i.c., target treatment or compared treatment) which is separated by the special [SE?] token. As shown in Fig. 4 the treatment "valsartan" " is appended to the original inputs to indicate that the patient is from the treat- ment group of "valsartan", The model processes the new inputs through the embedding layer and the Transformer encoder with parameters initialized with 0* We use the final hidden vector correspond- ing the first input token ( (CCLS) ) as the pooled representation [CLs)] from the pooling layer. We predict the potential outcomes under the treatment a via the linear head as f. o fe- (h reLs] (a)). The fine-tuning objective is the binary cross entropy (BCE) of the potential outcome prediction:
LrE (0' [= (BCE(ff fo- (hicus) (a)), vi]
Here, only the factual outcome are used for training loss computation as the counterfactual outcomes are unavailable in the observational data. After model fine-tuning, we infer the ITE d and ATE A as the difference between two predicted potential outcomes under the target and compared treatment:
ja= Target ÿa= C
E. ya Target
ÿa Compparrdd
(5)
4 EXPERIMENTS
In this section, we evaluate the proposed CURE from three aspects: 1) Quantitative analysis of the comparison performance with state-of-the-ar TEE methods on 4 downstream tasks; 2) Qualitative analysis including the validation of the estimated treatment effects with corresponding RCTs, and self-attention feature weights visualization; 3) Ablation studies including proposed feature embed- ding, pre- training data size, and generalizability of low-resource fine-tuning data.
Pre-training data. We extract patient data from MarketScan Commercial Claims and Encounters (CCAE) from 2012 to 2017, which contains andivduual-eved de-identified healthcare claims infor- mation from employers, health plans and hospitals. In this paper, we evolve patients who have ever been diagnosed with coronary artery disease (CAD) as our disease cohort. The definition of CAD is in Appendix C After conducting data preprocessing and study design, we obtain 2,955,399 patient sequences for pre-training. We obtain 9,435 medical codes including 282 diagnosis codes (i.e., we map the original ICD-9/10 billing codes into Clinical Classifications Software [CCS] and 9,153 medication codes (i.e., we map medications based on generic names from RED BOOK
Downstream tasks. As the ground truth treatment effects are not available in observational data, we use RCTs as the gold standard to verify our results. We focus on CAD-related RCTs which study the comparative effectiveness of two treatments for reducing the risk of stroke after CAD. We first collect all available Phase 2 and Phase 3 RCTs with CAD as disease name and stroke as disease outcome from commonly used as the primary outcome measurement in various CAD studies and it is well-defined in observational data. Then we select completed RCTs with published results. Finally, we end up with 4 RCTs that meet all the above criteria. We derive corresponding downstream tasks from our data based on the study design as specified in Fig. 2 More details of screening RCTs are in Appendix C
gov/ Stroke is selected because is
An additional semi-synthetic setting with comparison results in Appendix D
hhtps://www.ibmm http:/wwwww.ww//////om htwp:/ppooccc/mnnetttteeeeeeeeea
www.
www.hcup-us.ail htt.gggooooosstiwwwe/////ss/ss/

---

Under review as a conference paper at ICLR 2023
Baselines. We compare CURE with 8 neural network models for TEE, including state-of of -t-th--rrt methods. For models designed for continuous outcomes with mean square eπror (MSE) as a training objective, change the objective function to binary cross entropy for consistency. All the baselines are only trained on downstream data and are summarized below:
TARNet (Shalit al. 2017) predicts the potential outcomes based balanced representations
among treated and controlled groups.
2019) jointly optimizes treatment prediction and potential outcome pre- diction. The model first learns shared representations given the input data and then does prediction tasks via a three-head neural network: one for treatment prediction and two for potential outcomes.
DragonNet (Shi al.,
DR-CFR (Hassanpour & Greiner 2019) learns disentangled representations for counterfactual re- gression and assumes that the observed covariates can disentangled into three components: only contributing to treatment selection, only contributing to outcome predication, and both.
TNet (Curth & van der Schaar 202 is a neural network based T-learner (Künzel et al., 2019) (i.e.. type of meta-learners that decomposes the TEE into two sub-regression problems). TNet adopts two neural models as base learners for two potential outcomes.
SNer (Curth & van der Schaar 2021b) also learns disentangled representations and assumes that the observed covariates can be disentangled into five components by considering two potential out- comes separately.
FlexTENer (Curth & van der Schaar 202 incorporates the idea of inductive bias for the shared structure of two potential outcomes into TEE. The model adaptively learns what to share between the potential outcome functions.
Tr (Zhang et al.] 2022b) is a recently proposed Transformer-basee TEE model. The co- variates and treatments are encoded via a Transformer and aggregated for outcome prediction via cross-attention layer.
Base Model directly trains on the downstream datasets using the same architecture as CURE.
Metries. We evaluate the factual prediction performance using the standard classification metries:
Area under the ROC Curve (AUC) and Area under the Precision- RecciilonRRecall Curve (AUPR). We evaluate the counterfactual prediction performance using the influence function- heterogeneous effects (IF-PEHE) (Alaa & Van Der Schaar 2019), which helps to benchmark TEE methods when the ground truth effects are not available. Compared to the widely adopted precision of estimating heterogeneous effects (PEHE) that measures the mean squared error between estimated treatment effects and true treatment effects, IF-PEHE measures the mean squared error between estimated treatment effects and approximated true treatment effects. The output of the IF-PEHE metric is a numeric value and the lower the better. More details of this metric are in Appendix C
precision of estimating
Implementation details. Our pre- pre-training uses the BERT base architecture (Devlin et al.) 2018 with 768 hidden size, 12 attention heads, 12 layer Transformer and 3072 intermediate size. The maximum input sequence length is 256. The pre-training conducted on NVIDIA GeForce RTX 2080 Ti 1IGB GPUs with batch size of 96. We train our model using the adaptive moment estimation (Adam) optimizer, with an initial learning rate of le - 4 and learning rate warmup in the first 10% training steps. During the fine-tuning, the learning rate is 5e 5 without learning rate warmup. We fine-tune the model on each task for 2 epochs. The downstream data are randomly split into training. validation and test sets with percentages of 90%, 5%, 5% respectively. All results reported on the test sets. More implementation details are mentioned in Appendix C The code of our proposed CURE is available in Supplementary material.
4.1 QUANTITATIVE ANALYSIS
Comparison with sfate-of-the-art methods. Table shows the performance of factual outcome prediction (measured by AUC and AUPR) and TEE (measured by IF-PEHE) on four different down- stream tasks. We compare CURE with the state-of-the-ar TEE methods and report the results under 20 random runs. We observe that the proposed CURE has more than 3.8%, 6.9% and 15.7% respec- tive average AUC, AUPR and IF-PEHE improvement over the best baseline on these tasks. The results illustrate the promise and effectiveness of our proposed pre-training and fine-tuning method- ology for TEE. Notably, even without pre-training, the base model of CURE attains similar perfor- mance as the best baseline, which suggests the effectiveness of our architecture and data encoding designs. We demonstrate the model performance on addressing the treatment selection bias from quantitative and qualitative perspectives in Appendix D
7

---

Under review as a conference paper at ICLR 2023
Table 1: Comparison with state-of offte-of-fiee ahe methods on four downstream datasets. The results are the average and standard deviation over 20 runs.
Rivaroxahan v.s. Aspirin
Valsartan v.s. Ramipril AUPR 1
Method
AUC
AUPR
IF-PEHB 4
AUC f
[E-PEEE 4
TARNet
0.719 + 0.015
0.327 +: 0.023
0.546 0.044
0.683 +: 0.028
0.263 +: 0.029
0.545 4 0.061
DragonNet
0.767 +: 0.013
0.381 + 0.023
0.715 = 0.082
0.083 + 0.026
0.263 + 0.027
0.533 : 0.109
DR-CTR
0.759 +: 0.015
0.381 : 0.026
0.653 0.118
0.751 0.019
0.333 = 0.032
0.642 = 0.135
TNet
0.715 0.016
0.318 0.028
0.500 0.059
0.673 0.021
0.256 0.024
0.494 + 0.065
SNet
0.756 + 0.014
0,380 + 0.028
0.247 0.054
0.752 + 0.022
0.333 + 0.033
0.526 + 0.097
TEet
0.717 +: 0.014
0.319 + 0.022
0.565 0.054
0.662 +.002
0.241 + 0.027
0.602 + 0.088
TransTEE
0.717 +: 0.011
0.299 + 0.021
0.537
0.042
0.773
0.380 + 0.025
0.427 + 0.089
Base Model
0.758 + 0.029
0.406 + 0.037
0.180
0.038
0.780 + 0.050
0.305 + 0.06G
0.190 + 0.0G5
CURE
0.803 + 0.011
0.469 + 0.023
0.173 0.038
0.811 0.018
0.428 + 0.041
0.158 + 0.062
Ticagrelor v.s. Aspirin
Apixahan Warfarin
Method
AUC
AUPR
IF-PEHB 4
AUC 1
AUPR 1
IE-PEE /
TARNet
0.714 +: 0.008
0.359 + 0.016
0.520 0.048
0.748 + 0.012
0.447 +: 0.030
0.535 + 0.043
DragonNet
0.741 + 0.009
0.397 : 0.020
0.433 = 0.096
0.792 + 0.018
0.519 + 0.035
0.461 : 0.095
DR-CER
0.745 0.007
0.403 : 0.021
0.580 0.108
0.798 0.015
0.531 + 0.032
0.503 = 0.073
TNet
0.709 + 0.009
0.360 0.020
0.490 0.061
0.741 0.015
0.432 + 0.082
0.519 + 0.039
SNet
0.742 + 0.008
0.400 + 0.020
0.298 0.053
0.795 +0 0.014
0.525 + 0.034
0.414 + 0.054
FlexTENet
0.710 +: 0.010
0.351 +: 0.015
0.487 0.046
0.735 0.0022
0,413 + 0.033
0.578 + 0.030
TranSTEE
0.747 +: 0.022
0.385 : 0.015
0.387 0.021
0.799
0.517 = 0.031
0.409 + 0.059
Buse Model
0.751 #: 0.025
0.425 + 0.04
0.206
0.031
0.791 1 0.029
0.539 + 0.039
0.251 + 0.045
CURE
0.793 +: 0.008
0.489 + 0.024
0.198 0.068
0.826 + 0.014
0.588 + 0.024
0.224 + 0.066
4.2 QUALITATIVE ANALYSIS
Validate with RCT conclusion. As the ground truth treatment effects are not available in obser- vational data, we further evaluate the estimated treatment effects with corresponding ground truth RCTs. In Table 4.2, show the confidence intervals of estimated effects under 20 runs and RCT conclusions of each downstream task.
Table 2: Comparison of the estimated treatment effects with corresponding ground truth RCT. The estimated effects are shown in 95% confidence intervals (CI) under 20 bootstrap runs. The RCT conclusions are obtained from published articles.
Target v.s. Compared
Estimated Effect (CD)
Pvluu
Geuerated Hypothesis
Conclusion
Rivaroxaban v.s. Aspiria
[-0.009, 0.006]
0.452
No significant difference
No significatt diffcrance
Anand ct
2018
Valsartan Ramipril
1-0.003, 0.014]
0.103
No significunt differencee
No significant difference
Fieller 2021
Ticagrelor v.s. Aspirin
10.022. 0.040]
6c-14
T. is less effective than A.
No significant difference
Sandiny ct 2020
Apixahun v.s. Warfarin
[-0.039, 00002]
4e-4
A. more effective thun W.
A. more effective than W. Granger et al. 201
We use the direct difference to estimate the treatment effects (Hernán 2004). The results can be
interpreted as two potential conclusions: 1) The target treatment is sigmificantly more effective than the compared treatment in reducing the risk of the outcome if the upper bound of the confidence interval is lower than zero. (2) The target is not significantly more effective than the compared treatment if the confidence interval covers zero (i.e., no significant difference) or the lower bound is higher than zero (i.e., the compared treatment is more effective than the target treatment). As we can see, our estimated treatment effects are mostly consistent with each corresponding RCT conclusion. Though the generated hypothesis and RCT conclusion are not exactly the same for the third pair (Ticagrelor v.s. Aspirin), they both indicate that there is no significant reduced treatment effect of the target treatment over the compared treatment. The results demonstrate that our proposed CURE successfully identifies correct treatment effects using only observational patient data. The full results of all the baselines can be found in Appendix D
Self-attention visualization. The self-attention mechanism of the Transformer enables the explo- ration of interaction among input covariates and provides potential interpretation of the prediction results. We use Transformer visualization tool called bertviz (Vig (2019) to help visualize learned attention weights. We show the visualization results of some patient samples in Appendix
4.3 ABLATION STUDIES
Effect of embedding layer. We evaluate the effect of proposed time embedding (visit time and physical time) and type embedding respectively. As shown in Fig. [5] the model with both time and type embedding generally performs better than the other two embedding ablations. Especially,
8

---

Under review as a conference paper at ICLR 2023
incorporating time embedding yields larger performance improvement than the type embedding. This indicates that the proposed embedding method is better than the standard embedding method and time information plays more important role in TEE than the type information.
0.60
0.27
0.82-
0.55
0.24
0.81
2 0.80
0.50
0.21
0.45
0.18
0.15
R. v.S.A. A.
V.
u.. R. v.S.
T.
A. V.S.
R R.
v.S. A.
V.M.S.
T.V.S.
v.S. W
A
V.
v.s. R.
U.S.A
A. v.s.s W
w/o. time embedoling
w/o. type embedding
wi timeStype embedding
Figure 5: The effect of different embedding layer designs four downstream tasks.
Effeet of downstream data size. We demonstrate the model's 's effectiveness on the low resource of downstream data in Fig. 6 The plots show the model performance with different fractions of labeled downstream data. Generally, given only 5% 10% labeled data, the CURE achieves comparable per- formance to the Base Model which is trained on the fully labeled data. Specifically, the performance gains are large when given a small fraction of labeled data (1%-5%) and the curve tends to gently increase after the fraction is larger than 10%. With increased data size, the performance gradually fully labeled data. The results demonstrate that unsu- pervised pre-training benefits low- now-resource downstream tasks even when only a limited number of labeled data available for fine-tuning. More results of other metrics are in Appendix D
achieves the upper bound of fine-tuning
Rivaroxaban v.S. Aspirin
Valsartan Ramipril
Ticagrelor Aspirin
Apixaban v.s. Warfarin
0.80
0.80
0.80
0.75
0.75
0.75
AUC
AUC
0.70
0.70
0.65
0.65
0.70
10
20
0.60
5
10
20
5
10
20
10
20
Training Set Size (%)
Training Size
Training Set Size (%6
Training Set Size (%)
Low resource W. pre-training
Base Model
Full set w. pre-training
Figure 6: The effect of low resource in fine-tuning datasets four downstream tasks with different fractions of labeled training set (x-axes). The results are the average of 20 runs.
Effect of pre-training data size. We further explore the effect of pre-training data volume on the performance of downstream tasks. In Fig. 7 we show the AUC given different fractions of pre- training data. Here, 0% training set size denotes the Base Model, which is trained on the downstream data from scratch. Generally, the performance improves with the increase of pre-train data. The results indicate that pre-training is beneficial for downstream tasks by learning contextualized patient representations from large-scale unlabeled patient data. Other metrics are show in Appendix D
Rivaroxaban
Valsartan Ramipri
Ticagrelor Aspirin
Warfarin
0.81
0.53
0.5
0.50
$ 9.78
5 0.61
0.77
0.7
0.76
0.6
0.78
0.75
0.79
60
80
100
05
40
60
100
05
20
40
100
20
80
100
Taining
(54)
Training Set Size
Training Set (%)
Training
(5)
Figure 7: The effect of pre-training data volume on four downstream tasks (average of 20 runs).
CONCLUSION
In this paper, we study the problem TEE from observational data. We propose a new transformer- based TEE framework called CURE, which adopts the pre-training and fine-tuning paradigm. CURE is pre-trained on large-scale unlabeled patient data and then fine-tuned on labeled patient data for TER. We convert the structured patient data into sequence and design new sequence coding method to encode the structure and time into a comprehensive patient embedding. Thorough experiments show that pre-trale mre-training significantly boosts the TEE performance on 4 downstream tasks compared to state-of-the-ar methods. We further demonstrate the data scalability of CURE and ver- ify the results with corresponding published RCTs.

---

Under review as a conference paper at ICLR 2023
Ethics Statement. The observational data used in the paper from MarketScan Research Database, which is fully HIPAA-compliant de-identified, have very minimal risk of the potential for loss of privacy. Moreover, Per the DUAs with MarketScan, all users to access the data will need to take full research, ethics, and compliance training courses and be covered by IRBs. Thus, potential privacy and security risk would be eliminated and/or mitigated.
Reproducibility Statement. We provide the code and instructions needed to reproduce the results in supplemental material. The experimental data can be obtained from https: www. ibm. com/ products /tt.aaaaatteeeeee scan research databases
REFERENCES
Ahmed Alaa and Mihaela Van Der Schaar. Validating causal inference models via influence func- tions. In International Conference Machine Learning, pp. 191-201. PMLR, 2019.
Sonia S Anand, Jackie Bosch, John W Eikelboom, Stuart Connolly, Rafael Diaz, Peter Widimsky, Victor Aboyans, Marco Alings, Ajay Kakkar, Katalin Keltai, et al. Rivaroxaban with without aspirin in patients with stable peripheral or carotid artery disease: an international, randomised, double-blind, placebo-controlled trial. The Lancet, 391(10117):119-22,, 2018.
Tom Brown, Benjamin Mann, Nick Ryder. Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference on computer vision, pp. 219-229. Springer, 2020.
Zhengping Che, Sanjay Purushotham, Kyunghyun Cho, David Sontag, and Yan Liu. Recurrent neural networks for multivariate time series with missing values. Scientific reports, 8(1):1-12, 2018.
Tianqi Chen and Carlos Guestrin. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd aem sigkdd international conference on knowledge discovery and data mining, pp. 785-794, 2016.
Hugh A Chipman, Edward I George. and Robert E McCulloch. Bart: Bayesian additive regression trees. The Annals of Applied Statistics, 4(1):266-29)) 2988 2010.
Edward Choi, Mohammad Taha Bahadori, Andy Schuetz, Walter F Stewart, and Jimeng Sun. Doctor ai: Predicting clinical events via recurrent neural networks. In Machine learning for healthcare conference, pp. 301-318. PMLR, 2016a.
Edward Choi, Mohammad Taha Bahadori, Jimeng Sun, Joshua Kulas, Andy Schuetz, and Walter Stewart. Retain: An interpretable predictive model for healthcare using reverse time attention mechanism. Advances neural information processing systems, 29, 2016b.
Alicia Curth and Mihaela van der Schaar. On inductive biases for heterogeneous treatment effect estimation. Advances in Neural Information Processing Systems, 34, 2021a.
Alicia Curth and Mihaela van der Schaar. Nonparametric estimation of heterogeneous treatment effeects: From theory to learning algorithms. In International Conference on Artificial Intelligence and Statistics, pp. 18 1010-11 1818. PMLR, 2021b.
Rajeev H Dehejia and Sadek Wahba. Causal effects in nonexperimental studies: Reevaluating the evaluation of training programs. Journal of the American statistical Association, 94(448):1053- 1062, 1999.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: /8/0: 04.000 2018.
10

---

Under review as a conference paper at ICLR 2023
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. / 1929, 2020.
Thomas A Glass, Steven N Goodman, Miguel A Hernán, and Jonathan M Samet. Causal inference in public health. Annual review of public health, 34:61-75, 2013.
Christopher Granger, John H Alexander, John JV McMurray, Renato D Lopes, Elaine M Hylek, Michael Hanna, Hussein R Al-Khalidi, Jack Ansell, Dan Atar, Alvaro Avezum, et al. Apixaban versus warfarin in patients with atrial fibrillation. New England Journal of Medicine, 365(11); 981-992, 2011.
Zhenyu Guo, Shuai Zheng, Zhizhe Liu, Kun Yan, and Zhenfeng Zhu. Cetransformer: Casual ef- feet estimation via transformer based representation learning. In Chinese Conference on Pattern Recognition and Computer Vision (PRCV), pp. 524-535. Springer, 2021.
Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. In International Conference on Learning Representations, 2019.
Miguel Angel Hernán. A definition causal effect for cpidemiological research. Journal of Epi- demiology & Community Health, 58(4))265-217 27,, 2004.
Guido W Imbens. Nonparametric estimation of average treatment effects under exogeneity: A review. Review of Economics and statistics, 86(1):4-29, 2004.
Guido W Imbens and Donald B Rubin. Causal inference in statistics, social, and biomedical sci- ences. Cambridge University Press, 2015.
Sören R Künzel, Jasjeet Sekhon, Peter J Bickel, and Bin Yu. Metalearners for estimating heteroge- neous treatment effects using machine learning. Proceedings of the national academy of sciences, 116(10):41566 5664165, 2019.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. / 1692, 2019.
Junyu Luo, Muchao Ye, Cao Xiao, and Fenglong Ma. Hitanet: Hierarchical time-aware atten- tion networks for risk prediction on electronic health records. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery Data Mining, pp. 647-656, 2020.
Fenglong Ma, Radha Chitta, Jing Zhou, Quanzeng You, Tong Sun, and Jing Gao. Dipole: Di- agnosis prediction in healthcare via attention- bidirectional recurrent neural networks. In Proceedings of the 23rd ACM SIGKDD international conference on knowledge discovery and data mining, pp. 1903-1911, 2017.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning, pp. 4055- 4064. PMLR. 2018.
Marc A Pfeffer, Brian Claggett, Eldrin Lewis, Christopher B Granger, Lars Køber, Aldo Mag- gioni, Douglas L Mann, John JV McMurray, Jean-Lucien Rouleau, Scott D Solomon, et al. An- giotensin receptor- neprilysin inhibition in acute myocardial infarction. New England Journal of Medicine, 188-5208:18451855 2021.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. 2018
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, lly Sutskever, et al. Language models are unsupervised multitask leamners OpenAl blog. 1(8):9, 2019.
Laila Rasmy, Yang Xiang, Ziqian Xie, Cui Tao, and Degui Zhi. Med-bert: pretrained contextualized embeddings on large-scale structured electronic health records for disease prediction. NPJ digital medicine, 4(1):1-13, 2021.
11

---

Under review as a conference paper at ICLR 2023
Donald B Rubin. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Statistical Association, 100(469):322-933) 2005.
Sigrid E Sandner, Heribert Schunkert, Adnan Kastrati, Dominik Wiedemann, Martin Misfeld, An- dreas Boening, UIrich Tebbe, Bernd Nowak, Jan Stritzke, Guenther Laufer, et al. Ticagrelor monotherapy versus aspirin in patients undergoing multiple arterial or single arterial coronary artery bypass grafting: insights from the ticab trial. European Journal of -hoorcc Surgery, 57(4):732-739, 2020.
Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: gener- alization bounds and algorithms. In International Conference on Machine Learning, pp. 3076- 3085. PMLR. 2017.
Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. Advances in neural information processing systems, 32, 2019.
Yishai Shimoni, Ehud Karavani, Sivan Ravid, Peter Bak, Tan Hung Ng, Sharon Hensley Alford, Denise Meade, and Yaara Goldschmidt. An evaluation toolkit to guide model selection and cohort definition in causal inference. arXiv preprint arXiv: 190.:::990000444:: 2019.
John W Stanifer, Sean D Pokorney, Glenn M Chertow, Stefan H Hohnloser, Daniel M Wojdyla, Samira Garonzik, Wonkyung Byon, Ziad Hijazi, Renato Lopes, John H Alexander, et al. Apix- aban versus warfarin in patients with atrial fibrillation and advanced chronic kidney disease. Cir- culation, (11(11):13841111 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- tion processing systems, 30, 2017.
Jesse Vig. A multiscale visualization of attention in the transformer model. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics: System Demonstrations, pp. 37-42, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/ V1/711-3007. URL htthttps: www. www.wacllee web. org wttpr///nntoooloyy111-------0
Stefan Wager and Susan Athey. Estimation and inference of hetergeneous treatment effects using random forests. Journal of the American Statistical Association, 1133623):128--222,, 2018.
Linying Zhang, Yixin Wang, Martijn J Schuemie, David M Blei. and George Hripesak. Adjusting for indirectly measured confounding using large-scale propensity score. Journal of Biomedical Informatics, 134: 104204, 2022a.
Yi-Fam Zhang, Hanlin Zhang, Zachary C Lipton, Li Erran Li, and Eric P Xing. Can transformers be strong treatment effect estimators? arXiv preprint arXiv:2202.011 01336, 2022b.
12

---

Under review as a conference paper at ICLR 2023
APPENDIX
A CAUSAL ASSUMPTIONS
We follow the standard causal assumptions (Imbens & Rubin 2015 to help guarantee that the treat- ment effects are identifiable from the observational data.
Assumption (Consistency) The potential outcome under the treatment equals to the observed outcome if the actual treatments is a.
Assumption 2 (Positivity) Given the observational data of the history, if the the probability P(a- 1 a)) * then the probability of receiving treatment O or positive, i.e,, 0 < P(A a X a) < L, for all a E A and E X.
Assumption (Strong Ignorability) Given the observational data the history, the treatment as- signment is independent of the potential outcome, i.e., Y ( = a) 1 A x. for all E A.
Assumption T is fundamental to the potential outcome framework used to define counterfactuals and infer treatment effects. Essentially, this assumption requires that the treatment specified in the study must be precise enough that any variation within the treatment specification will not lead to a different outcome. 2 immples implies that all patients may receive the treatment whatever their observed covariates. Otherwise, is impossible to derive the counterfactuals for patients who do not have any chance of being in the other treatment group. Assumption 3 states that the poten- tial outcomes are independent of treatment assignment given the set of observed covariates. This assumption guarantees that the treatment effects are identifiable given the treatment, outcome and observed covariates as: E = 1) Y = 0)] = Ex x [ A = 1, x] - E| = 0, æ]].
Positivity in our data. We investigate the positivity assumption in each downstream data respec- tively. Following the positivity evaluation in Shimoni et al.] (2019), we estimate the propensity score of each individual via a propensity score model. The positivity can be evaluated through the distribution of propensity: whether the distribution is generally smooth with propensities normally distributed instead of much accumulation either propensity equals or propensity equals 1. As shown in Figure A1 the estimated propensity scores in each downstream dataset are generally nor- mally distributed without accumulation on either side. The empirical results demonstrate that our data generally satisfy the positivity assumption.
Rearrestan et Aagitn
Waltatton vr. Buniort
Tragreby oc. Aege'n
375
Ap Jaboe Werfern
200
Figure A1: Propensity distribution in each downstream dataset.
Strong ignorability in our data. The strong ignorability assumption should be generally valid in our data. First, the pre-training data have around 3M unlabeled patient sequences with 9,452 unique covariates. With such large patient data and rich covariate space, the model can identify sufficient potential confounders and effectively adjust the confounding bias. Moreover, recent study Zhang al. (2022a) of treatment effect estimation further demonstrates that the model trained on tens of thousands of covariates from observational data can even adjust for indirectly measured confounders. Second, according to our study design (as illustrated in Section 3.1 and Figure 2 all the covariates are obtained from the baseline period (time before the target treatment initiation) as potential confounders, and the outcomes are obtained from the follow-up period (time after the target treatment initiation). In this design, all the potential confounders collected prior to the first target treatment prescription (a.k.a,, pre-treatment covariates) thus simple collider bias or mediator bias will be mitigated.
13

---

Under review as a conference paper at ICLR 2023
B TRANSFORMER ARCHITECTURE
For each single Transformer encoder block, consists of multi-head self-attention layer followed by a fully-connected feed-forward layer (Vaswani 2017). The multi-head attention the most crucial part which can be calculated as,
MultiHead (h)
(oenct(heedd,... head,, WWo
head; Attention (hW? hw. hWY)
QKT )
Attention K, V) Softmax
vd
where h E Raxdmat denotes the hidden representations and d is the input sequence length. w.O RR xd w.K E Rdmmasd W.V Rddaaxdd Wo E Inndddammbar are learnable pa- rameter matrices. dmodsl/ and n the number of attention heads. We show the detailed model configuration in Fig.
"attention_pree htw..prrbs_dropout_.rooo 0.1, miesifiierdddooutt null, "hi thidden_act": "gelu", "hi dden_dropout_prohh 0.1,
"hi "hiddnnssiee": 768,
"ini "initiallier_rangee" 0.02,
"éntermedicte_size" 3072,
""ayer_nomm_ep". 1e-12,
Tmaxx_olyysicll_tineeeebbeee entiiidinmebeddnnng 13,
"max_pposttionn.mmbdddtnnngs"""" 512,
361,
"model_type": "bert", tmamaattention_hedda 12,
fmmmhhicden_layrr"". 12,
"pad_ "pad_token tokke_iddk 0, "pos wtw.iinnmmmodddigggtyyyoooe "absolute", "time_embest seneeeemmedddng true,
dtorr""dttpe:: "float32",
"transformers version": "4.17.0", sype__vccabssie""" 5, "use_ vuse_cache!; true, "vocob_size": 9452
Figure A2: Model configuration.
14

---

Under review as a conference paper at ICLR 2023
C ADDITIONAL DETAILS ON EXPERIMENTAL SETUPS
Pre-training data. The pre-training data are obtained from MarketScan Commercial Database [ which consists of medical and drug data from employers and health plans for over 215 million individuals. In this study, we focus on CAD as the studied disease and stroke as the outcome. The definitions of CAD and stroke are shown in Table Al and Table A2 respectively.
Table Al: The definition of coronary artery disease (CAD) from observational health data.
Reference (PMID)
16159046, 26524702, 28008010
A history of coronary revascularization in the EHR
Criteria
Or, history of acute coronary syndrome, ischemic heart disease, exertional angina
ICD-9 codes:
410* to 414*
ICD-10 codes:
The best approximation the following codes:
120 Angina pectoris
* Acute myocardial infarction
Diagnostic codes
122* Subsequent ST elevation (STEM)) and non-ST clevation (NSTEM) myocardial
infarction
123# Certain current complications following ST elevation (STEMI) and non-ST
elevation (NSTEMT) myocardial infarction (within the 28 day period)
124* Other acute ischemic heart diseases
125 Chronic ischemic heart disease
Table A2: The definition stroke from observational health data
Reference (PMID)
29202795
ICD-9 codes:
V12.54,
438.0-438.9
ICD 10 codes:
786.73
Diagnostic codes
160-169
subarachnoid hemorrhage (160);
intracerebral hemorrhage (I61);
cerebral infarction (163);
and other transient cerebral ischemic attacks and related syndromes and
transient cerebral ischemic attack (unspecified) (G458 and G459)
15

---

Under review as a conference paper at ICLR 2023
Downstream tasks. We demonstrate the flowchart for RCT extraction in Fig. A3 All RCTs extracted from with stroke the outcome and end up with 4 RCTs that satisfy all the above criteria. We have included all those RCTs for downstream task construction.
gov/l We start from 1,593 CAD-related RCTs
RCTs of CAD disease and Stroke as outcome
(n=1,593)
RCTs not in Phase 2 Phase (n=1,301)
RCTs Phase or Phase 3 (n#292)
RCTs not completed without results (n=221)
RCTs are completed with results
(n=71)
The studied drugs are neither valid available (n=67)
The studied drugs are valid and available in observational data (n=4)
Figure A3: The data flow for RCT extraction. The downstream tasks constructed based the extracted RCTs.
Table A3: The statistics of the downstream dataset.
Target Compared
Aspirin
Valsartan Ramipril
Ticagrelor v.s. Aqpirin
Apixahan v.s. Warfarin
of patients (Tunget: Compured) Female (%)
26340 (9569; 16771)
12850 (7306: 5544)
29248 (12477; 16771)
18187 (6701; 11486)
30.4
32.4
27.1
31.8
Age (group) index date
55-64
55-64
55-64
55-64
Paticnts stroke (%)
13.7
16.7
Average visits per patient
83.4
74.0
70.7
97.1
Average of codes per patient
182.3
157.2
152.0
215.9
Pre-training and fine-tuning data preparation. The pre-training is based on large-scale unla- beled patient data, and the fine-tuning based on small-scale labeled patient data, which are not used for pre-training. We first construct datasets for downstream tasks according the study de- sign of related randomized clinical trial (RCT). The patients who satisfy the eligibility criteria of the RCT included in each dataset respectively (see we construct the unlabeled pre-training data based on all the remaining patients who are not included in any of the downstream datasets. Therefore, the same patients do not appear simultaneously in both the pre-training stage and the fine-tuning stage.
details of the study design in Fig. (2) Then
Evaluation metrics. As the true treatment effects available in real-world data, the influence function-basee precision of estimating heterogeneous effects (IF-PEHE) (Alaa & Van Der Schaar 2019) for model evaluation. Following the same experimental setup, we calculate IF_ PEHE as,
Step Train two XGBoost (Chen & Guestrin 2016) classifiers for potential outcome prediction denoted by μo and P1, where Ha = Ply 1|X = x) using the training set Zurin- Then calculate the plug-in estimation T μr Nu Train XGBoost (Chen & Guestrin 2016) classifier propensity score function (i.e., the probability of receiving treatment) P(a 1|X x).
16

---

Under review as a conference paper at ICLR 2023
Given the estimated treatment effect ((x;) the test set Zest calculate the IF- PEHE with the influence function / as,
Step 2:
IF-PEHE = Σ [(T(x)) - 7(x))) I(x:))
2,E2m:
(2)
((x) =1 (1- + (y( -(x))- - -(x)) + (x) where W = (x)), B = 2a -
Implementation details. The pre-training model architecture follows the BERT base (Devlin al, 2018) and most hyperparameters remain the same default setting. The detailed hyperparameters setup is shown in Table folrr pre-training, and Table for fine-tuning. With NVIDIA GeForce RTX 2080 Ti 11GB GPUs, the pre-training takes about 20 hours with current setup. We have provided all code in supplemental material.
Table A4: Hyperparameters used in pre-training.
Parameters
CURE
Maximum Steps
100K
Initial Learning Rate
le-4
Batch Size
96
Warm-Up Steps
10K
Sequence Length
256
Dropout
0.1
Table A5: Hyperparameters search space and optimal parameters used for fine tie-tuning.
Parameters
Search Space
Optimal Value
Maximum Epochs
[1,2.3.4.5}
Initial Learning Rate
(1e-5, 3e-5, 5e-5)
5e-5
Batch Size
{16, 32, 64)
32
Sequence Length
256
256
Fixed Window Length
30
30
Baseline Window
(90, 180, 360, 720)
360
Dropout
0.1
0.1
Table A6: The parameter size of the proposed method and baselines.
Method
Model parameters
TARNet
2M
DragonNet
2M
DR-CFR
3M
TNet
4M
SNet
3M
Flex TexTENet
3M
Trans
7M
CURE
93M
17

---

Under review as a conference paper at ICLR 2023
Table A7: The influence of weight (a) associated with the discriminator in DragonNet the model performance on the Valsartan v.s. Ramipril dataset (random seed =42).
AUC
AUPR
IF-PEHE
0.2
0.677
0.304
0.768
0.4
0.679
0.308
0.689
0.680
0.310
0.660
0.8
0.682
0.312
0.644
1.0
0.683
0.314
0.643
1.2
0.679
0.315
0.593
1.4
0.681
0.317
0.584
1.6
0.682
0.318
0.589
1.8
0.683
0,319
0,586
2.0
0.685
0.321
0.595
CURE
0.805
0,428
0.161
18

---

Under review as a conference paper at ICLR 2023
D ADDITIONAL EXPERIMENTAL RESULTS
Visualization. The self-altention mechanism of the Transformer enables the exploration of inter- action among input covariates. The purpose of visualizing learned attention weights is to better understand our data and interpret what our model learns, how our model makes decisions, and in- vestigate whether the model learns the right associations or not.
As example, we show the attention weights of patient from Apixaban treatment group of Apix- aban v.s. Warfarin study in Fig. A4 Different colors denote the attention heads and there are 12 heads in total. The medications and diagnosis codes highlighted in the figure the most related features to the outcome prediction and treatment effect estimation. For example, Amiodarone is an antiarrhythmic medication used treat and prevent number of types of cardiac dysrhythmias including atrial fibrillation | A study (Stanifer et al. 2020) shows that apixaban is superior to war- farin in preventing stroke patients with atrial fibrillation. Those attention weights could be used to analyze the treatment effects in some subgroups that characterized by the attended feature set.
meducationnAtorassttin.CCaalll
medicationnLLisnoprii mediccttoon-Metorool__rrtraea
medication.Furoosmiiie
medication:Diggxin
diagnosis:Nonspeciffi chest pain
diagnosis:Acute and unspecified renal failure diagnosis:Fluid and electrolyte disorders diagnosis:Cardiac dysrhythmias
Figure A4: The visualization of Top 10 attention weights associated with the special token [CLS] of patient from Apixaban treatment group.
Ablation studies in effect of embedding layer. We note that in Figure type embedding is helpful to out of 4 datasets as excluding type embedding (w/o. type embedding) yields an increase in IF-PEHE scores in datasets (R v.s. A, V v.s. R and T v.s. A) and drop in IF-PEHE score dataset (A v.s. W). According to the basic statistics of 4 datasets (Table A3). we find that the average number of visits per patient and the average number of codes per patient the largest in A v.s. W dataset among all the 4 datasets. In this case, the type embedding for the data with a larger number of visits/codes may be helpful in the data with a smaller number of visits/codes. And this characteristic is captured by the IF-PEHE metric. Oppositely, excluding type embedding in T v.s. A (which contains the smallest number visits/codes per patient) yields a significant performance drop (i.e.. type embedding is important to T v.s. A dataset). The results demonstrate that each embedding plays different role different datasets.
Ablation studies effect of downstream data size. As shown in Fig. A5 our model can achieve comparable performance as measured by AUPR the Base Model with only around 50%% la- beled downstream data.
Ablation studies in effect of pre-training data size. As shown in Fig. A6 the model gradually yields better performance in terms of AUPR scores when the pre-train data size increases.
Evaluation on non-random assignment. We show the t-SNE visualization of learned patient representations of treatment and control group respectively (see Fig. A7 The visualization of our model (CURE) demonstrates that the distribution variance between two groups is marginal and the non-random assignment issue is alleviated.
httpps:/ /ttp:://wwww.. http://wwww..g....e// drugs. wtwpggg..oo//oooogrrggaeeeaeeeee hwtp://aaaaaarro...eoooom
19

---

Under review as a conference paper at ICLR 2023
Rivaroxaban Aspirin
Valsartan Ramipril
Ticagrelor Aspirin
Apixaban v.5. Warfarin
0.60
0.40
0.45
0.55
0.35
AURR
0.50
0.30
0.25
0.35
0.40
0.35
0.20
0.30
10
20
5
10
20
5
10
20
10
20
Trainingg Set Size (%)
Trairing Set Size (%)
Training Set Size 1%6)
Training Set Size (%)
Low resource W. pre-training
Base Model
Full set w. pre-training
Figure A5: The effect of low resource in fine-tuning datasets on four downstream tasks with different fractions of labeled training (x-axes). The results the average of 20
Rivaroxaban Aspirin
Valsartan Ramipri
Ticagrelor Aspirin
Apixaban Warfarin
0.46
0.42
0.48
3.50
g 0.45
g 9.41
0.47
0.44
0.43
0.42
R
0.39
0.38
0.40
0.45
0.44
0:46
0.55
8:37
80
100
05
40
60
05
20
40
100
20
80
100
Taaninng
Training
Training Set
Training
(56)
Figure A6: The effect of pre-training data volume on four downstream tasks (average of 20 runs).
We further adapt the design in DragonNet (Shi et al. 2019) in our model's fine-tuning stage. Specif- ically, we add an additional prediction head for propensity score estimation and modify the loss function incorporate both outcome prediction and propensity prediction. We compare the new model (CURE+propensityy with the proposed CURE model on downstream tasks respectively. As shown in Table A8 the performance of these two models is comparable in terms of both factual prediction and treatment effect estimation.
Valsartan v.s. Ramipril (DR-CFR)
Valsartan v.s. Ramipril (CURE)
treated
20
0
-10
-20
control
-30-
-40
-40
-20
20
40
-40
-20
0
20
40
Figure A7: The t-SNE visualization of learned patient representations on Valsartan v.s. Ramipril dataset.
20

---

Under review as a conference paper at ICLR 2023
Table A8: Comparison of the proposed model (CURE) and new model based on CURE but adapt- ing propensity module in DragonNet (Shi et al. 2019) (CURE + propensity) on 4 downstream tasks (random 42ed=42).
Rivaroxaban v.s. Aspirin
Valsartan Ramipril
Method
AUC
AUPR
IF-PEHE
AUC
AUPR
IF-PEHE
CURE
0.786
0.419
0.186
0.805
0.428
0.161
CURE+propensity
0.789
0.427
0.178
0.804
0.428
0.156
Ticagrelor Aspirin
Apixaban Warfarin
Method
AUC
AUPR
IF-PEHE
AUC
AUPR
IF-PEHE
CURE
0.807
0.502
0.211
0.838
0.597
0.182
CURE-propensity
0.806
0.505
0.239
0.844
0.602
0.253
21

---

Under review as a conference paper at ICLR 2023
Evaluation with RCT conclusion. We use the bootstrap percentile method which approximates the data sample statistic from the distribution of the bootstrap sample statistic. Specifically, for each downstream data, first obtain the estimated (individual) treatment effects under 100 ran- dom runs with different random seeds. We then calculate the average treatment effects (c) of each bootstrap and sort them from smallest to biggest. Finally the (95%) CIs are estimated as Cls = (Co.05/2, C100.5/2) where Co denotes the a x 100 percentile of the sorted list.
We evaluate the estimated treatment effects on 100 bootstrap samples. As shown in Table D the results are consistent with the reported results on 20 bootstrap samples: only a little change in numerical values but no influence on the generated hypothesis.
Table A9; Comparison of the estimated treatment effects with corresponding ground truth RCT. The estimated effects are shown in 95% confidence intervals (CI) under 100 bootstrap runs, The RCT conclusions are obtained from published articles.
Target v.s. Compared
Estimated Effect (CI)
valae
Generated Hypothesis
RCT Coaclusion
Rivaroxaban v.s. Aspirin
[-0.009, 0.008)
0.4140
No significant difference
No significant difference Anand et 2018
Valsartan v.s. Ramipeil
1-0.004, 0.011]
0.0582
No significant difference
No significant differencd Peti a ST
Ticugrelor Aspicin
10.019. 0.043]
5.7814e-25
T. is less effective than A.
No sigmificans difference
Apixahan v.s.
[-0.037. 00002]
1.32576-15
A. more effective than
is more effective W Grargggrr
We evaluate the treatment effects estimated by all the baselines and conduct the same hypothesis testing. As shown in Table (oble/AT0bellxs our method correctly generates 3 (out of 4) RCT conclusions that match the ground truth RCT conclusions while the best baselines only identify 2 (out of 4) RCT conclusions.
Table A10: Comparison of the estimated treatment effects with corresponding ground truth RCT of all methods.
Rivaroxabun v.s. Aspicin
Valsartan Ramipril
Method
Estimated Effect (CT)
valuu
Match RCT Conclusion?
Estimated Effect (CT)
vvause
Match RCT Cooclusion?
TARNet
[0.066, 0.095]
5.678c-10
No
[-0.357 -0.003)
0.026
DragonNet
10.18, 0,256]
5,979c-12
No
J0.03, 0.07]
4.681c-05
No
DR CTR
10.13, 0813
2.783e-10
No
10,002, 0.04]
0.033
No
TNet
[0.041, 0.07)
2.509e-07
No
[-0.038, -0.001]
0.039
No
SNel
[-0,002, 0,008]
0.231
Yes
[-0.051, -0.026]
3.168c-06
No
FlexTENet
[0.064, 0.108]
1.529c-07
No
[-0.079, -0.035]
3.184e-05
No
TransTEE
f-0.013, -0.0021
0.018
No
[-0.019.
0.420
Yos
CURE
[-0.009, 0,006]
0.452
Yes
[-0.003, 0.014)
0,103
Yes
Ticagreler v.s. Aspitin
Apixaban v.s. Warfarin
Method
Estimated Effect (CD)
vvluue
Match RCT Conclusion?
Estimatedd Effect (CD
P valuu
Match RCT Conclusion?
TARNet
[0.064, 0.101]
2.861c-08
No
[-0.006, 0.028]
0.207
No
DragonNet
[-0.013, 0.01]
0,821
Yes
[0.018, 0.056]
6.281c-04
No
DR CFR
F-0066,, -0.029]
4.915c-05
No
[-0.026, -0.002)
0.047
Yes
TNet
[0.046, 0.069]
6.474c-09
No
[0.009, 0.0231
2.329e-04
No
SNet
[0.005, 0.016]
4,398c-04
No
[-0.046 -0.017]
2.112a-04
Yes
FlexTENet
[0.045,0.068]
5,243c-09
No
[0.012, 0,042]
0.001
No
TransTEE
[-0.014, -0.009
0.0216
No
[-0.027, -0.002]
0.027
Yos
CURE
[0.022, 0.040]
5.982 14
No
[-0.039, .0.002)
4e-04
Yes
Semi-synthetic experiment. We generate a semi-synthetic dataset based on real patient data ob- tained from the MarketScan data. Specifically, we simulate treatment assignment and potential outcome using pre-treatment covariates z (i.e., historical co-medication, co-morbidities and de- mographies). The treatment assignment is simulated by a|z ~ Bernoull(Sigmoid((" + m)), where s /N & denotes the aggregation of all historical covariates. The outcome is simulated by y/a, a ~ Bernoulli (Sigooid( (w= Ba + n)), where w ~ Nool' 0.. I). 3 ~ NO, 1), n ~ (O,00.1)
1 0.1 I). [V] is the cardinality of medical feature vocabulary,
(O,00..))
As we have all potential outcomes under both treatment and control arms available in the semi- synthetic data, the model performance is evaluated with Precision of Estimating Heterogeneous Effects (PEHF), which measures the root mean square error between the true treatment effect and
22

---

Under review as a conference paper at ICLR 2023
estimated treatment effect. The comparison results from the semi-synthetic dataset shown in Table [II]] The proposed model CURE yields the best performance among all the baselines. We will add the results for the semi-synthetic dataset in revision.
Table All: Comparison with oftte-of-thh-arr methods semi-synthetic Valsartan dataset. The results are the average and standard deviation over 20 runs.
Ramipril
Method
PEHE
TARNet
0.768 0.012
DragonNet
0.759 + 0.015
DR-CFR
0.714 0.014
TNet
0.784 = 0.017
SNet
0.776 0.022
FlexTENet
0.791 = 0.014
TranSTEE
0.689 + 0.012
CURE
0.596 + 0.010
23

---

Under review as a conference paper at ICLR 2023
Additional related work in clinical risk prediction. Though some existing works Choi et al. (2016b/a): Ma et al. (2017): Che et al.] (2018); Luo et al. (2020) in clinical risk prediction using electronic medical records (EHRs) also have patient encoding before sending the imput into the model, the patient encoding is unique our problem scenario and more comprehensive than the existing patient encoding in clinical risk prediction.
Fundamentally, clinical risk prediction is very different from treatment effect estimation in prob- lem formulation, model design, and experiment setup. To the best of our knowledge, we the first study to incorporate this comprehensive patient encoding with the following pre-training and fine-tuning framework as whole representation learning patient data for treatment effect es- timation problems. Even if ignoring the great differences between these two tasks, our patient data encoding is not exactly the same as the encoding methods in existing work clinical risk prediction (as shown in Table A12) Specifically, besides the token embedding, the proposed patient encoding considers 1) physical time information, which denotes the time gap between the observation and in- dex date, can capture the (irregular) time information: 2) visit time information is crucial to identify the relationship between each medical code and corresponding visit time; 3) type information to help capture the heterogeneity of the medical data (i.e., medications, diagnosis codes and demographics).
Table A12: Patient encoding methods in the clinical risk prediction literature. Time interval between two observations. ** Position embedding in Transformer Vaswani et al. (2017)
Backhone
tnken (variahle) emb.
physical time emb.
visit time emb.
type emb.
RETAIN Choi et al.
2016b
RNN
Yes
No Yes*
No
No
Doctor Al hor cl al 12016
RNN RNN
Yes
No
Dipole Ma
2017
Yes
No
No
No
GRU-D el ul. 2018
RNN
Yes
Yes*
No
No
HiTANet Luo et al 2020
Transformer
Yes
Yes
Yes"a
No
Mod-BERT [assy ct 2021
Transformer
Yes
No
CURE (Ours)
Transformer
Yes
Yes
Yes Yes
To further verify that our proposed embedding gives better performance, We compare our method with GRU-D Che et al.) (2018) which is one of the most representative works in clinical risk pre- diction that explicitly considers the time information. The patient encoding in GRU-D uses the time interval between two observations information. We implement the patient encoding method in GRU-D in our model by integrating the token embedding and time interval embedding the patient embedding and keeping the remaining model architecture the same. We conduct comparison experiments on 4 downstream data respec- tively. As shown in Table A13 our proposed patient coding method has more than 1.1%, 2.8% and 2.4% respective average AUC, AUPR and IF-PEHE improvement over the patient encoding method in GRU-D.
the input to the GRU unit in order to capture the time
Compared to the patient encoding method in GRU-D, our proposed patient encoding method is more comprehensive. For time information, we consider both visit time information and physical time information. Here, the physical time information can capture the (irregular) time information and plays a similar role as the time interval embedding in GRU-D. The visit time information, which not considered in GRU-D, is crucial to identify the relationship between each medical code and the corresponding visit. Besides the time information, we additionally consider the type information to help better capture the heterogeneity of the medical data.
24

---

Under review as a conference paper at ICLR 2023
Table A13: Comparison of the patient encoding in CURE with the patient encoding in GRU-D Che
al. (2018) 4 downstream tasks.
Rivaroxaban v.s. Aspirin
Valsartan v.s. Ramipril
Method AUC
AUPR
IF-PEHE
AUC
AUPR
IF-PEHE
GRU-D 0.791
0,448
0.179
0.802
0.395
0.176
CURE
0.803
0.469
0.173
0.811
0.428
0.158
Ticagrelor v.s. Aspirin
Apixaban Warfarin
Method
AUC
AUPR
IF-PEHE
AUC
AUPR 0.561
IF-PEHE
GRU-D
0.784 0.793
0.458
0.244
0.813
0.248
CURE
0.489
0.198
0.826
0.588
0.224
25


Reviews:
[{"review_id": "h0SDcAHK11", "reviewer": "Reviewer_9BnT", "summary": "", "questions": "", "limitations": "", "rating": 5, "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 paper tackles the problem of the estimation of the effect of a treatment with specific reference to the problem of comparing different treatments on given outcomes. The author/s describe/s a framework obtained by combining transformer-based pre-training with fine-tuning by calling it CURE, which is specifically designed for observational data. In particular, the CURE framework is first trained on unlabled data with the aim to obtain a contextual representation of patients, then a fine-tuning step exploits the few available labeled data for comparing the effect of two treatments on a given outcome of interest. The main contributions of the paper, in my understanding, are the sequential encoding devoted to represent longitudinal patient data, an embedding that the paper states to incorporates both structure and time. A second contribution is represened by a set of numerical experiments, togheter with their associated results that, according to what the author/s state/s also share the interesting property of data scalability.", "summary_of_the_review": "I found the paper to be not easy to read and how it is presented does not make it easy to appreciate the main contributions of the paper itself. The proposed framework, namely CURE, is definitely new, or at least I'm not aware of any other similar framework that has been presented in the specialized literature. However, the paper lacks in clarity and the presented case study has criticalities in terms of the fact that health care data suffer from missingness and several biases, while apparently the paper does not discuss in detail these fundamenatl aspects of the consider domain. Furthermore, also the notation is difficult to follow and in some cases it seem to be a mismatch, e.g. the argument of function f in formula (3). However, results of numerical experiments witness in favour of CURE when compared to state of the art approaches.", "strength_and_weaknesses": "Strength\n- the problem tackled is a relevant and urgent one\n- results of numerical experiments seems to witness the proposed approach compares favourably to state of the art methods\n- the idea of pre-training using unlabled data is useful, even if in a general context it is not new.\n\nWeaknesses\n- the issue of missing data is not taken into account, while this is a major one in the case where observational data are taken into account in the health care domain\n- some aspects about the assumptions and type of issues that are taken into account are not clear, i.e., I read \"The model trained for one particular problem or data may fail to generalize to other scenarios.\", what do you precisely mean by generalization? interval validity, external validity ?\n- the assumption of positivity is a strong one and I wonder whether it can hold in the practical domain presented by the paper, at least I would like to have seen in the results that the author/s take into account and discuss on this aspect.\n- at page 4 when defining x it is not clear which comorbidities are taken into account, those at the beginning of the longitudinal sample or those that have developed through time?\n- the assumption that all covariates are confounders or potential confounders is weak and averaging with respect to all covariates, expecially post-treatment covariates is a very bas practice that can greatly affect the performance to compare causal effect of different treatments.\n- formula (2) seems to implicitly assume independence between token, type, visit and physical, I think more comments are needed here to justify\n- m at the end of page 5 seems to indicate a different quantity tha what presented in previous pages, this make hard to understand the paper.\n- I miss the concept of \"causal treatment\", I know about treatment and causal effect of a treatment but I ignore what a \"causal treatment\" is, would you be that kind to explain?\n- apparently the paper states to be in the conditions to compute counterfactual, this is theoretically impossible unless you provide a theoretical proof of what you state. It is clear I could have missinterpreted what you wrote and if this is the case simply reply by clarifying that you did not intended to say that you can compute counterfactuals with your model. Eventually you can estimate counterfactural, provided a set of assumptions and that you show how this goal is achieved fro a theoretical point of view.\n- how to manage the fact that the study population is different from the target population? I think the paper must comment on this fundamental issue\n- the results are not discussed in a clear manner and thus not potential justitification or explanation are given to help the interested reader to appreciate the paper\n- I read \"interaction among input covariates\" but I miss what you mean by this, correlation? any other type of association? However, they are not enough to explain anything about the causal story behind the data\n- in formula (3) with respect to the fucntion f, you use the argument x while few lines before you used the argument e, thus I'm truly puzzled beacause x and e are defined as different quantities in a previous part of the paper.", "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": "_Gjrk6R4QtP", "reviewer": "Reviewer_BLi7", "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": "The authors propose a transformer-based pre-training and fine-tuning framework for TEE from observational data. The proposed method is pre-trained to learn representative contextual patient representations. Then, fine-tuned on labeled patient data for TEE. The proposed method is evaluated on 4 downstream TEE tasks.\n\n", "summary_of_the_review": "The presentation of the paper is clear. It is valuable to study motivation. But technological innovation is limited.", "strength_and_weaknesses": "Strengths:\n\n*The authors provided an abundantly detailed ablation study of the different embeddings used and their effect on the performance of the model.\n\n*The article is well organized. Detailed experiments and analysis. \n\nWeaknesses:\n\n* Given that the benefits of the standard pre-training / fine-tuning approach are well-known, the results in this paper are not particularly novel or surprising.\n\n* Technological innovation is insufficient.", "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": "_96BB8qqMy3", "reviewer": "Reviewer_w381", "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 proposes a model and pre-training method for estimating conditional expected outcomes as part of causal estimation and evaluates this approach using real and semi-synthetic data.", "summary_of_the_review": "I am of two minds on this paper. On the one hand, I do not think the proposed method is particularly novel and I think the paper is missing a major chunk of literature on clinical risk prediction. On the other hand, I was very impressed by the experimental setup and comparison to real clinical trials, which I found very thoughtfully done. As this is a machine learning conference and I think the machine learning portion of the paper is lacking, I recommend a marginal rejection.\n\nMajor comments:\n1. Many, many papers have been written about clinical risk prediction using EHR data, including several that explicitly consider how to encode temporal information (e.g., Che et al. (2018)). I do not think the authors can make the claim that they have developed \"a new patient data encoding method to encode structured observational patient data and incorporate covariate type and time into patient embeddings.\" without reviewing this literature.\n\n2. I do not consider Section 3.1 to be a scientific contribution. These are the preprocessing steps that every EHR-based study must perform and I do not see anything particularly unique about the proposed steps. Please correct me if I have missed something.\n\n3. Notation in Section 2, specifically $y_1(x) - y_0(x)$ is not precisely defined.\n\n4. It is good that the authors mention causal identifiability assumptions in Section 2, but these assumptions cannot be made universally. Rather, for each causal analysis, the authors need to consider the various sources of confounding and censoring and make an argument that there are sufficient observed variables to account for potential biases. No such argument is made for the cases considered. \n\n5. In Section 5.2: 20 bootstrap samples is not enough to reliably estimate 95% confidence intervals. I recommend at least 100, but more would be better. Additionally, there are several ways estimate CIs from bootstrap samples. What method was used here?", "strength_and_weaknesses": "See below.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "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": "Irk1SLjLfDj", "reviewer": "Reviewer_wo54", "summary": "", "questions": "", "limitations": "", "rating": 8, "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 proposes to use a large dataset to train a pre-trained BERT-like deep model for patient representation using unsupervised learning approach. It then fine-tine the pre-trained model using data with labels for Treatment Effect Estimation. The motivation is based on the fact that the data with labels is limited. The proposed learning methodology CURE is evaluated based on a carefully prepared claims data and RCT data for treatment effect estimation. Performance improvement is observed based on the proposed CURE, and ablation study is also performed to account for the effectiveness of different tricks used in the proposed CURE framework.", "summary_of_the_review": "This paper demonstrates the effectiveness of fine-turning a pre-trained model trained using a large unlabelled data for treatment effect estimation. Comprehensive evaluation using a large scale dataset has demonstrated the potential impact of the proposed methodology to be put into practice.", "strength_and_weaknesses": "Strengths:\n+ Demonstrate the effectiveness of fine-tuning a pre-trained model for patient representation to gain performance improvement\n+ Datasets are carefully prepared and the experiments are carefully designed\n+ A lot of detailed implementation details are provided in the Appendix\n+ BERT is used as the basic architecture, which has been widely adopted and the proposed methodology should be applicable in practice.\n\nWeaknesses:\n- Some particular design of time embedding and type embedding is introduced. It is not clear how optimal the design is, and whether there can be alternatives.  \n- Not much discussion on the issue of selection bias possibly caused by confounder which is important for treatment effect estimation. The use of propensity network is mentioned in the appendix but not in the main text. It seems good to have some related discussion in the main text.\n\nQuestions:\n- It seems that the results of time embedding only were not reported in Figure 5. \n- In term of IF-PEHE, the type embedding seems hurting for three cases and useful for one case (A v.s. W). Any reason? \n\n", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject