SlowGuess/NeurIPS_2023_OCR huggingface dataset:
C.1 Experimental setup
Here, we provide more details about the experimental setup for each investigated setting. This serves to complement the high-level overview given in Table 1. Experiments were run using Google Cloud Services. Deep learning-based methods (i.e., CaML and its ablations, S-learner w/ meta-learning, T-learner w/ meta-learning, SIN, GraphITE, FlexTENET, TARNet, and DragonNet) were run on n1-highmem-64 machines with 4x NVIDIA T4 GPU devices. The remaining baselines (RA-learner, R-learner, X-learner, and T-learner) were run on n1-highmem-64 machines featuring 64 CPUs.
Fair comparison. We perform hyper-parameter optimization with random search for all models, with the meta-testing dataset predetermined and held out. To avoid "hyperparameter hacking", hyperparameters ranges are consistent between methods wherever possible, and were chosen using defaults similar to prior work [35, 25]. Choice of final model hyper-parameters was determined using performance metrics (specific to each dataset) computed on the meta-validation dataset, using the best hyper-parameters over 48 runs (6 servers x 4 NVIDIA T4 GPUs per server x 2 runs per GPU) (Appendix C.4). All table results are computed as the mean across 8 runs of the final model with distinct random seeds.
C.1.1 Claims dataset
Interventions (W): We consider drug prescriptions consisting of either one drug, or two drugs prescribed in combination. We observed 745 unique single drugs, and 22,883 unique drug pairs, excluding interventions which occurred less than 500 times. Time of intervention corresponds to the first day of exposure. To obtain intervention information, we generated pre-trained drug embeddings from a large-scale biomedical knowledge graph [8] (see Appendix C.5). Drugs correspond to nodes in the knowledge graph, which are linked to other nodes (e.g. genes, based on the protein target of the drug). Drug combination embeddings are the sum of the embeddings for their constituent drugs.
Control group. A challenge in such causal analyses of clinical settings is defining a control group. We randomly sample  $5\%$  (1.52M patients) to use as controls, with a 40/20/40 split between meta-train/meta-val/meta-test. When sampling a natural experiment for a given intervention, we select all patients from this control group that did not receive such an intervention. An additional challenge is defining time of intervention for the control group. It is not possible to naively sample a random date, because there are large quiet periods in the claims dataset in which no data is logged. We thus sample a date in which the control patient received a random drug, and thus our measure of CATE estimates the increase in side effect likelihood from the drug(s)  $W$ , compared to another drug intervention chosen at random.
Outcome  $(Y)$ : We focus on the side effect pancytopenia: a deficiency across all three blood cell lines (red blood cells, white blood cells, and platelets). Pancytopenia is life-threatening, with a  $10-20\%$  mortality rate [38, 43], and is a rare side effect of many common medications [42] (e.g. arthritis and cancer drugs), which in turn require intensive monitoring of the blood work. Our outcome is defined as the (binary) occurrence of pancytopenia within 90 days of intervention exposure.
Features  $(X)$ : Following prior work [24], patient medical history features were constructed by time-binned counts of each unique medical code (diagnosis, procedure, lab result, drug prescription) before the drug was prescribed. In total, 443,940 features were generated from the following time bins: 0-24 hours, 24-48 hours, 2-7 days, 8-30 days, and 31-90 days, 91-365 days, and  $365+$  days prior. All individuals in the dataset provided by the insurance company had at least 50 unique days of claims data.
Metrics: We rely on best practices for evaluating CATE estimators as established established by recent work [86, 11], which recommend to assess treatment rules by comparing subgroups across different quantiles of estimated CATE. We follow the high vs. others RATE (rank-weighted average treatment effect) approach from Yadlowsky et. al [86], which computes the difference in average treatment effect (ATE) of the top  $u$  percent of individuals (ranked by predicted CATE), versus all individuals:
$$
R A T E @ u = \mathbb {E} \left[ Y (1) - Y (0) \mid F _ {S} (S (X)) \geq 1 - u \right] - \mathbb {E} \left[ Y (1) - Y (0) \right], \tag {14}
$$
where  $S(\cdot)$  is a priority score which ranks samples lowest to highest predicted CATE, and  $F_{S}(\cdot)$  is the cumulative distribution function (CDF) of  $S(X_{i})$ . For instance, RATE @ 0.99 would be the difference between the top  $1\%$  of the samples (by estimated CATE) vs. the average treatment effect (ATE) across all samples, which we would expect to be high if the CATE estimator is accurate. The real-world use case of our model would be preventing drug prescription a small subset of high-risk individuals. Thus, more specifically, for each task  $j$ , intervention  $w_{j}$  in the meta-dataset, and meta-model  $\Psi_{\theta}$  (our priority score  $S(\cdot)$ ), we compute  $RATE@u$  for each  $u$  in [0.999, 0.998, 0.995, 0.99] across individuals who received the intervention.
We now summarize how to estimate RATE performance metrics for a single intervention (task). As RATE performance is calculated separately per-intervention we are concerned with a single intervention, we use the simplified notation (i.e.  $Y_{i}(1)$  instead of  $Y_{i}(w)$ ) from Section 3. Due to the fundamental problem of causal inference (we can only observe  $Y_{i}(0)$  or  $Y_{i}(1)$  for a given sample), the true RATE, as defined above, cannot be directly observed.
We follow the method outlined in Section 2.2 and 2.4 of Yadlowsky et. al, [86] in which we compute  $\widehat{\Gamma}_i$ , a (noisy but unbiased) estimate for CATE which is in turn used to estimate RATE:
$$
\mathbb {E} \left[ \widehat {\Gamma} _ {i} \mid X _ {i} \right] \approx \tau \left(X _ {i}\right) = \mathbb {E} \left[ Y _ {i} (1) - Y _ {i} (0) \mid X _ {i} \right]. \tag {15}
$$
Our data is observational, and as such we can estimate  $\widehat{\Gamma}_i$  using a direct non-parametric estimator [80]:
$$
\widehat {\Gamma} _ {i} = W _ {i} \left(Y _ {i} - \hat {m} \left(X _ {i}, 0\right)\right) + \left(1 - W _ {i}\right) \left(\hat {m} \left(X _ {i}, 1\right) - Y _ {i}\right) \tag {16}
$$
$$
m (x, w) = \mathbb {E} \left[ Y _ {i} (w) \mid X _ {i} = x \right] \tag {17}
$$
where  $m(x, w)$  is a model that predicts the outcome. Here  $\hat{m}(x, w)$  represent nonparametric estimates of  $m(x, w)$ , respectively, which we obtain by fitting a cross-fitting a model to the intervention natural experiment over 5-folds. We use random forest models for  $\hat{m}(x, w)$ , as they perform well (achieving  $\geq 0.90$  ROC AUC across all meta-testing tasks for predicting outcomes) and are robust to choice of hyperparameters.
RATE can then be estimated via sample-averaging estimator. Specifically, we compute the difference between the average value of  $\widehat{\Gamma}_i$  for those in the top  $u$  percent of individuals (based on our metamodel's predictions), compared to the average  $\widehat{\Gamma}_i$  across all individuals. For further discussion on estimating RATE, we refer readers to [86]. Note that estimates of RATE are unbounded: RATE can be less than 0 (due to predictions inversely relating to CATE).
Finally, because our meta-testing dataset consists of individuals treated with drugs known in the medical literature to cause pancytopenia (identified by filtering drugs using the side effect database SIDER [42]), observational metrics of recall and precision are also a rough proxy for successful CATE estimation. Thus, as secondary metrics, we also compute Recall @ u and Precision @ u for the same set of thresholds as RATE, where a positive label is defined as occurrence of pancytopenia after intervention. We find that these metrics are highly correlated to RATE in our performance results.
Training & Evaluation: For each method, we ran a hyperparameter search with 48 random configurations (48 due to running 8 jobs in parallel on 6 servers each) that were drawn uniformly from a pre-defined hyperparameter search space (see Appendix C.4). Methods that can be trained on multiple tasks to then be applied to tasks unseen during training (i.e., CaML and its ablations, S-learner w/ meta-learning, T-learner w/ meta-learning, SIN, GraphITE) were trained for 24 hours (per run) on the meta-training tasks. Model selection was performed on the meta-validation tasks by maximizing the mean RATE@0.998 across meta-validation tasks. Then, the best hyperparameter configuration was used to fit 8 repetition runs across 8 different random seeds. Each repetition model was then tested on the meta-testing tasks, where for all metrics averages across the testing tasks are reported. To make the setting of multi-task models comparable with single-task models that were trained on meta-testing tasks (requiring a train and test split of each meta-testing task), the evaluation of all models was computed on the test split of the meta-testing tasks, respectively. Single-task baselines (FlexTENET, TARNet, and DragonNet, RA-learner, R-learner, X-learner, and T-learner) were given access to the meta-testing tasks during training. Specifically, model selection was performed on the meta-validation tasks, while the best hyperparameter configuration was used to train 8 repetition models (using 8 random seeds) on the train split of each meta-testing task. For the final evaluation, each single-task model that was fit on meta-testing task  $i$  was tested on the test split of the same meta-testing task  $i$ , and the average metrics were reported across meta-testing tasks.
C.1.2 LINCS
Interventions (W): Interventions in the LINCS dataset consist of a single perturbagen (small molecule). For intervention information, we used the molecular embeddings for each perturbagen using the RDKit featurizer. The same cell line-perturbagen combinations are tested with different perturbagen dosages and times of exposure. [46]. To maintain consistency in experimental conditions while also ensuring that the dataset is sufficiently large for training a model, we filter for most frequently occurring dosage and time of exposure in the dataset, which are  $10\mu M$  and 24 hours, respectively. We use data from 10,322 different perturbagens.
Control group. For each perturbagen (at a given timepoint and dose), we use cell lines which did not receive that intervention as the control group.
Outcomes  $(Y)$ : We measure gene expression across the top-50 and top-20 landmark differentially expressed genes (DEGs) in the LINCS dataset. Accurately predicting in gene expression in these DEGs is most crucial to the drug discovery process.
Features  $(X)$ : We use 19,221 features from the Cancer Cell Line Encyclopedia (CCLE) [22] to describe each cell-line, based on historical gene expression values in a different lab environment. Our dataset consisted of 99 unique cell lines (after filtering for cell-lines with CCLE features).
Metrics: A key advantage of experiments on cells is that at evaluation time we can observe both  $Y(0)$  and  $Y(1)$  for the same cell line  $X$ , through multiple experiments on clones of the same cell line in controlled lab conditions. In the LINCS dataset,  $Y(0)$  is also measured for all cells which received an intervention. Thus, we can directly compute the Precision Estimation of Heterogenous Effects (PEHE) on all treated cells in our meta-testing dataset. PEHE is a standard metric for CATE estimation performance [30], analogous to mean squared error (MSE).
$$
P E H E = \frac {1}{N} \sum_ {i = 1} ^ {N} \left(\tau_ {i} - \hat {\tau} _ {i}\right) ^ {2} \tag {18}
$$
Training & Evaluation: For each method, we ran a hyperparameter search with 48 random configurations (48 due to running 8 jobs in parallel on 6 servers each) that were drawn uniformly from a pre-defined hyperparameter search space (see Appendix C.4). Methods that can be trained on multiple tasks to then be applied to tasks unseen during training (i.e., CaML and its ablations, S-learner w/ meta-learning, T-learner w/ meta-learning, SIN) were trained for 12 hours (per run) on the meta-training tasks. Model selection was performed on the meta-validation tasks by minimizing the overall PEHE for the Top-20 most differentially expressed genes (DEGs) across meta-validation tasks. Then, the best hyperparameter configuration was used to fit 8 repetition runs across 8 different random seeds. Each repetition model was then tested on the meta-testing tasks, where for all metrics averages across the testing tasks are reported.
Data augmentation: We augment each batch of data during training to also include treated samples that have their pseudo-outcome labels to 0, and their  $W$  set to the zero vector.
C.2 Selecting holdout interventions for meta-validation and meta-testing
C.2.1 Claims.
In the 30.4 million patient insurance claims dataset, each intervention task in meta-train/meta-val/meta-testing corresponds to a natural experiment of multiple patients, with some interventions (e.g. commonly prescribed drugs) having millions of associated patients who were prescribed the drug. One challenge is that in this setting, there is overlap in subjects between the natural experiments sampled by CaML, which can lead to data leakage between training and testing. For instance, if a patient received Drug 1 (in meta-test) and Drug 2 (meta-train), they would appear in both natural experiments, resulting in data leakage.
We take a conservative approach and exclude all patients who have ever received a meta-testing drug in their lifespan from the natural experiments for meta-val/meta-train. Similarly, we exclude all patients who received a meta-validation drug from meta-training.
This approach means we must take great care in selecting meta-testing drugs. Specifically, we must trade off between selecting drugs that are important (covering enough patients) while not diminishing the training dataset size. For instance selecting a commonly prescribed (e.g. aspirin) for meta-testing would deplete our meta-training dataset by over  $50\%$  of patients. Thus we only selected meta-test/meta-validation drugs which were prescribed to between 1,000,000 and 100K patients in our dataset, after filtering for only drugs which known to cause Pancytopenia [42] (using the SIDER database). From this subset of drugs, we randomly selected 10 meta-testing drugs and 2 meta-validation drugs, resulting in a total meta-testing/meta-validation pool of 4.1 million patients and 685K patients respectively.
To evaluate on unseen pairs of drugs on the same hold-out test dataset, we additionally created a second pairs testing dataset from the 5 most frequently occurring combinations from the meta-testing dataset. This allowed us to train a single model on the same meta-train split and evaluate on both single drug and drug pair interventions without occurrence of data leakage. Designing a larger evaluation of pairs was not possible because while pairs of drugs are commonly prescribed as intervention, each particular pair of drugs is a rare event, and accurately evaluating CATE estimation performance (for a rare outcome such as Pancytopenia) requires amassing a natural experiment with at least several thousand patients who received the same intervention.
C.2.2 LINCS.
The goal in selecting holdout interventions for the meta-validation and meta-testing sets was to ensure that they consisted of both cell lines and tasks (small molecules) that had not been seen previously at the time of training (i.e. zero-shot on cell lines and tasks).
Using a random data splitting approach would result in large portions (up to  $50\%$ ) of the data being unused to comply with the zero-shot requirements on cell lines and tasks. One approach to tackle this was to reserve only those tasks in the held-out sets which had been tested on the fewest cell lines. This preserved the maximum amount of data but resulted in an average of just 1 cell line per task in the meta-testing and meta-validation sets, which would not be fair to the non-zero shot baselines.
To address these issues, we designed a new data split procedure that exploits the structure of how tasks and cell lines are paired. To do so, We first clustered tasks by the cell lines they are tested on. We then identified a set of 600 drugs that had all been tested on a shared set of roughly 20 cell lines. We divided the cell lines and tasks within this set into the meta-validation and meta-testing set, while enforcing zero-shot constraints on both. This resulted in roughly 10 cell lines per intervention in both the meta-validation and meta-testing sets, while still maintaining a reasonably large size of 11 distinct cell lines and 300 distinct tasks in both sets. All remaining tasks and cell lines were reserved for the training set. (See Table 8)
C.3 Understanding CaML's performance
Our comparison to CATE estimators which are restricted to single interventions (Grey, Table 2,5) shows that a key reason for CaML's strong performance is the ability to jointly learn across from many intervention datasets, in order to generalize to unseen intervention.
Additionally, in both the Claims and LINCS settings, we conduct two key ablation studies to further understand the underlying reason for CaML's strong performance results.
In our first ablation experiment (w/o meta-learning), we trained the CaML model without employing meta-learning, instead using the standard empirical risk minimization (ERM) technique [78]. This can be seen as a specific implementation of the CaML algorithm (refer to Algorithm 1) when  $k = 1$  [54]. The results of this experiment showed a varying degree of performance deterioration across our primary tests. In the Claims settings, we observed a decrease in the RATE performance metric by  $15\% -22\%$  (refer to Table 2), while in the LINCS settings, the PEHE performance metric decreased by approximately 0.01 (see Table 3). These results indicate that the absence of meta-learning affects the model's performance, although the impact varies depending on the specific setting. An important detail to consider is that the Claims data experiments dealt with substantially larger datasets, each comprising hundreds of thousands of patients per intervention. This extensive scale of data potentially amplifies the benefits of using meta-learning in the CaML model for the Claims dataset. The larger dataset enables the model to adapt to a given task over a larger set of iterations without reusing the same data, thereby enhancing the efficacy of meta-learning.
Our second ablation (w/o RA-learner) assesses the sensitivity of CaML's performance to different pseudo-outcome estimation strategies. A key aspect of CaML is flexibility in choice of any pseudo-outcome estimator to infer CATE, in contrast to prior work which uses specific CATE estimation strategies [25, 35]. We find that CaML performance benefits strongly from flexibility of pseudo-outcome estimator choice. We assess this by using an alternative pseudo-outcome estimator. Firstly, we find that this ablation results in much noisier model training. For instance, the standard deviation in RATE across the 8 random seeds increases by  $20\times$  when using the alternative pseudo-outcome estimator in the claims setting. Moreover, the alternative pseudo-outcome estimator typically worsens performance, decreasing RATE by up to  $6\%$  in the Claims setting, and increasing PEHE by  $20\% -21\%$  in the LINCS setting (Table 3). We note that this ablation performs slightly better at the 0.99 threshold, which may be a result of the high variance in this ablation. Specific choice of alternative pseudo-outcome estimator for this ablation varies by setting. We use the R-learner [55] for Claims as it also achieves strong single task performance (Table 2, grey) on Claims data. However, R-learner is restricted to single-dimensional outcomes, and thus for LINCS (in which outcomes are 50 and 20 dimensional), we use the PW-learner instead [16].
C.4 Hyperparameter space
C.4.1 Claims dataset hyperparameter space
We list the hyperparameter search spaces for the medical claims dataset in the following tables. Table 9 represents the search space for CaML. The SIN baseline consists of two stages, Stage 1 and Stage 2. For the Stage 1 model, we searched the identical hyperparameter search space as for CaML (Table 9). For Stage 2, we used the hyperparameters displayed in Table 10. The search space for the GraphITE baseline is displayed in Table 11. For the S-learner and T-learner w/ meta-learning baselines, we use the same hyperparameter space as for CaML (Table 9) with the only major difference that the these baselines predicts the outcome  $Y$  instead of  $\hat{\tau}$ . For all deep learning-based methods, we employed a batch size of 8,192, except for GraphITE, where we were restricted to using a batch size of 512 due to larger memory requirements. Single-task neural network baselines (FlexTENet, TARNet, and DragonNet) are shown in Tables 12,13, and 14, respectively. For the remaining baselines, i.e., the model-agnostic CATE estimators, the (shared) hyperparameter search space is shown in Table 15. Finally, applied L1 regularization to the encoder layer of the customizable neural network models (that were not reused as external packages), i.e., SIN learner, GraphITE, T-learner w/ meta-learning, and S-learner w/ meta-learning, and CaML.
C.4.2 LINCS hyperparameter space
We list the hyperparameter search spaces for LINCS in the following tables. CaMLis shown in Table 16. SIN Stage 1 used the same search space as CaML (Table 16. The search space of SIN Stage 2 is shown in Table 17. S learner and T-learner w/ meta-learning used the same search space as CaML. The search space of GraphITE is shown in Table 18. All methods that were applied to LINCS used a batch size of 20.
C.5 More details on intervention information
Here we give more details about the intervention information used for the medical claims dataset. In order to perform zero-shot generalization, we acquired information about a specific intervention through the use of pretrained embeddings. We generated these embeddings on the Precision Medicine Knowledge Graph [8] that contains drug nodes as well as 9 other node types. We extracted embeddings for 7957 drugs from the knowledge graph.
To extract rich neighborhood information from the knowledge graph we used Stargraph [49], which is a coarse-to-fine representation learning algorithm. StarGraph generates a subgraph for each node by sampling from its neighbor nodes (all nodes in the one-hop neighborhood) and anchor nodes (a preselected subset of nodes appearing in the multihop neighborhood). In our case the anchor nodes were the  $2\%$  of graph nodes with the highest degree. For the scoring function we used the augmented version of TripleRE [90] presented in the StarGraph article [49].
We performed a hyperparameter optimization to compare different models and determine the one we used to calculate our final embeddings (see Table C.5). The hyperparameter search was random with the objective of minimizing the loss function used in training on held out data. The search range for each of the parameters is displayed in C.5. Since certain parameters did not seem to influence the final score as much we decided to use them as constants and focus on optimizing the hyperparameters in the table. Therefore the number of sampled anchors was set to 20 and  $u = 0.1$  in the augmented TripleRE function, the values matching those seen in Stargraph [48].
Our final embeddings were 256-dimensional, the learning rate was 2e-4, the drop-ratio was 5e-3. We used the self-adversarial negative sampling loss with  $\gamma = 8$  and we sampled 4 neighbor nodes for each subgraph.