SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We formulate the setting as a supervised learning problem with noisy labels (pseudo-outcomes) where we learn a smooth function  $f = \Psi(w, x) \rightarrow \tau$  among a family  $\mathcal{F}$ . We focus on  $\tau \in \mathbb{R}$ , and assume  $\tau \in [0,1]$  without loss of generality, since we can normalize  $\tau$  to this range. The training dataset has  $n$  interventions with  $m$  samples each, i.e., first  $n$  i.i.d. draws from  $P_W$ :  $w^{(1)}, \ldots, w^{(n)}$  and then for each  $w^{(j)}$ ,  $m$  i.i.d. draws from  $P_X$ :  $x_1^{(j)}, \ldots, x_m^{(j)}$ .
The main theorem quantifies the rate that combining information across different interventions helps with zero-shot performance. We prove a finite-sample generalization bound for the ERM variant of CaML. The ERM is a special case of ADAPT with  $k = 1$  that is more conducive to rigorous analysis. The advantage of Reptile over ERM is orthogonal and we refer the readers to the original discussion [54]. We assume the estimated pseudo-outcomes  $\tilde{\tau}$  during training satisfy  $\tilde{\tau} = \tau + \xi$  where  $\xi$  is an independent zero-mean noise with  $|\xi| \leq \epsilon$  almost surely for some  $\epsilon \geq 0$ ,
$$
\hat {f} = \min  _ {f \in \mathcal {F}} \hat {L} (f) = \min  _ {f} \frac {1}{n m} \sum_ {j = 1} ^ {n} \sum_ {i = 1} ^ {m} \left(f \left(w ^ {(j)}, x _ {i} ^ {(j)}\right) - \tilde {\tau} _ {i} ^ {(j)}\right) ^ {2}.
$$
The test error is  $L(f) = \mathbb{E}_{W,X,\tau}[(f(w,x) - \tau)^2]$ . Let  $f^{*} = \min_{f}L(f)$ . We bound the excess loss  $L(\hat{f}) - L(f^{*})$ . Our key assumption is that interventions with similar features  $W$  have similar effects in expectation. We assume that all functions in our family are smooth with respect to  $W$ , i.e.,  $\forall f\in \mathcal{F},\mathbb{E}_{W,X}\left[\| \partial f / \partial W\| _2^2\right]\leq \beta^2$ .
Theorem 1. Under our assumptions, with probability  $1 - \delta$
$$
\begin{array}{l} L (\hat {f}) \leq L (f ^ {*}) + 8 (1 + \epsilon) R _ {n m} (\mathcal {F}) + 8 \sqrt {\frac {(1 + \epsilon) R _ {n m} (\mathcal {F}) \log (1 / \delta)}{n}} + \frac {2 \log (1 / \delta)}{3 n} + \\ (1 + \epsilon) \sqrt {\frac {(3 2 C \beta^ {2} + 2 (1 + \epsilon) ^ {2} / m) \log (1 / \delta)}{n}} \\ \end{array}
$$
where  $R_{nm}$  is a novel notion of zero-shot Rademacher complexity defined in equation (9);  $C$  is a Poincaré constant that only depends on the distribution of  $W$ . For large  $n, m$ , the leading terms are the function complexity  $R_{nm}(\mathcal{F})$ , and an  $O(\sqrt{1/n})$  term with a numerator that scales with  $\beta$  and  $(1 + \epsilon)^2/m$ . This validates our intuition that when the intervention information  $W$  is more informative of the true treatment effects (smaller  $\beta$ ), and when the estimation of  $\tau$  in the training dataset is more accurate, the performance is better on novel interventions. Please refer to Section A for the full proof. Compared to standard generalization bound which usually has a  $\sqrt{1/n}$  term, our main technical innovation involves bounding the variance by the smoothness of the function class plus Poincaré-type inequalities. When  $\beta$  is much smaller than 1 we achieve a tighter bound.
Dataset	Samples	Features (X)	Outcome (Y)	Intervention type	Intervention information (W)
Claims	Patients	Patient history (binned counts of medical codes)	Pancytopenia onset	Drug intake (pre-scription)	Drug embedding (knowledge graph)
LINCS	Cell lines	Cancer cell encyclope-dia	Expression of land-mark genes (DEG)	Perturbagen (small molecule)	Molecular embeddings (RDKit)
Table 1: High-level overview of our two experimental settings. Details in Appendix C.1.
We explore to what extent zero-shot generalization is practical when predicting the effects of interventions. We thus design two novel evaluation settings using real-world data in domains where zero-shot CATE estimation will be highly impactful: (1) Health Insurance Claims: predicting the effect of a drug on a patient, and (2) LINCS: predicting the effect of a perturbation on a cell. We use new datasets because existing causal inference benchmarks [31, 73] focus on a single intervention. By contrast, zero-shot causal learning must be conceptualized in a multi-intervention setting.
Zero-shot Evaluation. Each task corresponds to estimating CATE for a single intervention, across many individual samples (e.g. patients). We split all tasks into meta-training/meta-validation, and a hold-out meta-testing set for evaluating zero-shot predictions (Table 2, unseen drugs for Claims and Table 3, unseen molecular perturbations in LINCS). For the Claims dataset, we also consider the challenging setting of combinations of unseen drugs (Table 5).
Each meta-validation and meta-testing task contains a natural experiment of many samples (e.g., patients) who received the unseen intervention, and many control samples who did not receive the intervention. The same patient (Claims) or cell-line (LINCS) can appear in multiple tasks (if they received different interventions at different times). Thus, to ensure a fair zero-shot evaluation, we exclude all samples who have ever received a meta-testing intervention from meta-val/meta-train. Similarly, we exclude all meta-validation patients from meta-train. Details on holdout selection are provided in Appendix C.2.
Table 1 gives an overview of both benchmarks. In the Claims dataset, we compare zero-shot predictions with strong single-intervention baselines which cannot generalize to unseen interventions. To do so, we further split each task in meta-validation and meta-testing into a train/test (50/50) split of samples. These baselines are trained on a task's train split, and all methods are evaluated on the test split of the meta-testing tasks. On the LINCS dataset, as each task consists of  $< 100$  cells, single-intervention baselines performed weakly and are excluded from analysis.
Baselines. We compare the zero-shot performance of CaML to two distinct categories of baselines. (1) Trained directly on test interventions. These are strong CATE estimators from prior work and
can only be trained on a single intervention. Thus, we train a single model on each meta-testing task's train split, and evaluate performance on its test split. This category includes T-learner [44], X-learner [44], RA-learner [16], R-learner [55], DragonNet [72], TARNet [70], and FlexTENet [17].
(2) Zero-shot baselines are trained across all meta-training tasks and are able to incorporate intervention information  $(W)$ . These methods are thus, in principle, capable of generalizing to unseen interventions. We use GraphITE [25] and Structured Intervention Networks (SIN) [35]. We also introduce two strong baselines which learn to directly estimate  $Y(w)$  and  $Y(0)$  by meta-learning across all training interventions, without using pseudo-outcomes: S-learner and T-learner with meta-learning. These extend the S-learner and T-learner from prior work [44] to incorporate intervention information  $(W)$  in their predictions. We elaborate on implementation details of baselines in Appendix C.7. For details on hyperparameter search and fair comparison, see Appendix C.1.
Ablations. In our first ablation experiment (w/o meta-learning), we trained the CaML model without meta-learning, instead using the standard empirical risk minimization (ERM) technique [78]. Our second ablation (w/o RA-learner) assesses the sensitivity of CaML's performance to different pseudooutcome estimation strategies. For further details on how these ablation studies were implemented, see Appendix C.3. We discuss the key findings from these ablations in Section 6.3.
6.1 Setting 1: Personalized drug side effect prediction from large-scale medical claims
Our first setting (Claims) is to predict the increased likelihood of a life-threatening side effect caused by a drug prescription. We leverage a large-scale insurance claims dataset of over 3.5 billion claims across 30.6 million patients in the United States<sup>2</sup>. Each date-stamped insurance claim contains a set of diagnoses (ICD-10 codes), drug prescriptions (DrugBank ID), procedures (ICD-10 codes), and laboratory results (LOINC codes). Laboratory results were categorized by whether the result was high, low, normal, abnormal (for non-continuous labs), or unknown.
Interventions are administration of one drug  $(n = 745)$ , or two drugs  $(n = 22,883)$  prescribed in combination. Time of intervention corresponds to the first day of exposure. Intervention information  $(W)$  was generated from pre-trained drug embeddings from a large-scale biomedical knowledge graph [8] (Appendix C). We compute drug combination embeddings as the sum of the embeddings of the constituent drugs. We focus on the binary outcome  $(Y)$  of the occurrence of the side effect pancytopenia within 90 days of intervention exposure. Pancytopenia is 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. Following prior work [24], patient medical history features  $(X)$  were constructed by time-binned counts of each unique medical code (diagnosis, procedure, lab result, drug prescription) at seven different time scales before the drug was prescribed, resulting in a total of 443,940 features. For more details, refer to Appendix C.1.
Metrics We rely on best practices for evaluating CATE estimators in observational data, as 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 (for more details, see Appendix C.1). For instance, RATE @ 0.99 is 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. Note that estimates of RATE can be negative if model predictions are inversely associated with CATE. We elaborate on the RATE computation in Appendix C.1.
The real-world use case of our model is preventing drug prescription for a small subset of high-risk individuals. Thus, more specifically, for each task  $j$ , intervention  $w_{j}$  in the meta-dataset, and metamodel  $\Psi_{\theta}$ , we compute  $RATE@u$  for each  $u$  in [0.999, 0.998, 0.995, 0.99] across individuals who received the intervention. We use a narrow range for  $u$  because pancytopenia is a very rare event occurring in less than  $0.3\%$  of the patients in our dataset. Hence, in a real-world deployment scenario, it is necessary to isolate the small subset of high-risk patients from the vast majority of patients for whom there is no risk of pancytopenia onset.
Insurance company undisclosed per data use agreement.
RATE @u(↑)	Recall @u(↑)	Precision @u(↑)
0.999	.998	0.995	0.99	0.999	0.998	0.995	0.99	0.999	0.998	0.995	0.99
Random	0.00	0.00	0.00	0.00	0.00	0.00	0.01	0.00	0.00	0.00	0.00	0.00
T-learner	0.32	0.26	0.16	0.10	0.12	0.18	0.26	0.31	0.36	0.29	0.18	0.11
X-learner	0.06	0.05	0.04	0.03	0.02	0.04	0.08	0.12	0.09	0.07	0.06	0.05
R-learner	0.19	0.17	0.12	0.08	0.06	0.1	0.19	0.26	0.24	0.21	0.15	0.11
RA-learner	0.47	0.37	0.23	0.14	0.17	0.26	0.38	0.45	0.54	0.42	0.26	0.16
DragonNet	0.09	0.07	0.05	0.04	0.03	0.05	0.08	0.11	0.15	0.12	0.08	0.06
TARNet	0.15	0.12	0.07	0.05	0.05	0.08	0.12	0.14	0.18	0.15	0.09	0.06
FlexTENet	0.10	0.09	0.06	0.04	0.04	0.06	0.11	0.16	0.15	0.13	0.09	0.06
GraphITE	0.19	0.12	0.05	0.03	0.07	0.08	0.09	0.10	0.23	0.14	0.07	0.04
SIN	0.00	0.00	0.00	0.00	0.00	0.00	0.01	0.02	0.01	0.01	0.01	0.01
S-learner w/ meta-learning	0.21	0.16	0.09	0.05	0.08	0.11	0.15	0.16	0.25	0.18	0.1	0.06
T-learner w/ meta-learning	0.40	0.31	0.18	0.11	0.15	0.22	0.32	0.38	0.45	0.35	0.21	0.13
CaML - w/o meta-learning	0.39	0.31	0.18	0.11	0.15	0.22	0.32	0.39	0.45	0.35	0.22	0.14
CaML - w/o RA-learner	0.45	0.36	0.22	0.14	0.16	0.24	0.34	0.41	0.48	0.38	0.26	0.16
CaML (ours)	0.48	0.38	0.23	0.13	0.18	0.27	0.38	0.45	0.54	0.43	0.26	0.16
Table 2: Performance results for the Claims dataset (predicting the effect of drug exposure on pancytopenia onset from patient medical history). Key findings are (1) CaML outperforms all zero-shot baselines (RATE is  $18 - 27\%$  higher than T-Learner w/ meta-learning, the strongest zero-shot baseline) (2) CaML performs stronger (up to  $8\times$  higher RATE values) than 6 of the 7 baselines which are trained directly on the test interventions, and performs comparably to the strongest baseline trained directly on the test interventions (RA-learner). Mean is reported across all runs; standard deviations included in (Appendix Table 4). Analogous trends hold for generalization to pairs of unseen drugs (Appendix Table 5).
Additionally, because our meta-testing dataset consists of individuals treated with drugs known to cause pancytopenia, observational metrics of recall and precision are also a rough proxy for successful CATE estimation (and highly correlated to RATE, Table 2). 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.
6.2 Setting 2: Cellular gene expression response due to perturbation
Our second setting (LINCS) is to predict how a cell's gene expression  $(Y)$  will respond to intervention from perturbagen (small molecule compound such as a drug). This is a critical problem as accurately predicting intervention response will accelerate drug-discovery. We use data for 10,325 different perturbagens from the LINCS Program [74]. Each perturbagen corresponds to a different small molecule. Molecular embeddings were generated using the RDKit featurizer [46] and used as intervention information  $(W)$ . Outcomes  $(Y)$  of interest are post-intervention gene expression across the top-50 and top-20 differentially expressed landmark genes (DEGs) in the LINCS dataset. We did not look at all 978 genes since most do not show significant variation upon perturbation. We use 19,221 features  $(X)$  from the Cancer Cell Line Encyclopedia (CCLE) [22] to characterize each cell-line  $(n = 99)$ , each of which correspond to unperturbed gene expression measured in a different lab environment using a different experimental assay. For more details, see Appendix C.1.
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 in estimating heterogeneous effects (PEHE) on all treated cells in our meta-testing dataset, an established measure for CATE estimation performance analogous to mean-squared error [30] (see Appendix C.1).
6.3 Key findings
CaML's zero-shot predictions outperform baselines with direct access to the target intervention. In the medical claims setting, single intervention baselines (Tables 2, dark grey rows) are the highest performing baselines as we train them directly on the meta-test intervention. Still, CaML outperforms 6 out of 7 of these baselines (up to  $8 \times$  higher RATE) and achieves comparable performance to the strongest of these baselines, the RA-learner. Furthermore, CaML strongly outperforms alternative zero-shot CATE estimators (RATE is  $18 - 27\%$  higher than T-Learner w/ meta-learning, the strongest zero-shot baseline). In the LINCS data, multi-intervention learners are strongest as
PEHE 50 DEGs  $(\downarrow)$  PEHE 20 DEGs  $(\downarrow)$
Mean.	3.78	4.11
GraphITE	3.58 ± 0.023	3.82 ± 0.011
SIN	3.78± 0.001	4.06 ± 0.001
S-learner w/ meta-learning	3.63 ± 0.004	3.90 ± 0.004
T-learner w/ meta-learning	3.61 ± 0.007	3.85 ± 0.006
CaML - w/o meta-learning	3.57 ± 0.006	3.79 ± 0.004
CaML - w/o RA-learner	4.28 ± 0.517	4.60 ± 0.413
CaML (ours)	3.56 ± 0.001	3.78 ± 0.005
Table 3: Performance results for the LINCS dataset (predicting the effect of an unseen perturbation on the gene expression of an unseen cell-line). CaML outperforms all baselines. Improvement is largest for the 20 most differentially expressed genes, where most signal is expected.
there are only a small number of instances (cell lines) per intervention<sup>3</sup>. CaML outperforms both single-intervention and multi-intervention learners by drawing from both of their strengths—it allows us to use strong CATE estimation methods (i.e. the RA-learner) which previously were restricted to single interventions, while sharing information across multiple interventions.
CaML learns to generalize from single interventions to combinations of unseen interventions (drug pairs). We evaluate CaML's performance in the challenging setting of predicting the personalized effects of combinations of two drugs which are both unseen during training, while only training on interventions consisting of single drugs. CaML achieves strong performance results (see Appendix Table 5), surpassing the best baseline trained on the test tasks, and outperforms all zero-shot baselines, across all 12 metrics.
Understanding CaML's performance results. Our ablation studies explain that CaML's performance gains are due to (1) our meta-learning formulation and algorithm (in contrast to the w/o meta-learning row, in which ERM is used to train the model), and (2) the flexible CATE estimation strategy, allowing to take advantage of recently developed CATE estimators previously restricted to single interventions (in contrast to the w/o RA-learner row, in which an alternative pseudo-outcome estimator is used). Lastly, (3) comparison to existing binary intervention CATE estimators trained separately on each meta-testing intervention (Table 2, grey rows) shows that we gain from learning from thousands of interventions. See Appendix C.3 for details on ablations.
We introduce a novel approach to predict the effects of novel interventions. CaML consistently outperforms state-of-the-art baselines, by unlocking zero-shot capability for many recently developed CATE estimation methods which were previously restricted to studying single interventions in isolation. While our study is limited to retrospective data, we plan to prospectively validate our findings. Future work includes designing new model architectures and CATE estimators which learn well under the CaML framework, developing new metrics to evaluate zero-shot CATE estimators, as well as more generally exploring novel learning strategies that enable zero-shot causal learning.
Societal impacts. In high-stakes decision-making inaccurate predictions can lead to severe consequences. It is important not to overly rely on model predictions and proactively involve domain experts, such as doctors, in the decision-making process. Additionally, it is crucial to ensure that underserved communities are not disadvantaged by errors in treatment effect estimates due to underrepresentation in the training data. Important avenues for achieving equitable CATE estimation in future work include process-oriented approaches (i.e., evaluating model errors for underserved demographics), and outcome-oriented methods (i.e., gauging model impacts on demographic utility) [12, 57, 69, 2]. Furthermore, the deployment of CATE models could raise privacy concerns. These models typically require access to individual patient data to estimate personalized treatment effects accurately. Ensuring the privacy and security of this sensitive information is crucial to avoid potential data breaches or unauthorized access, which could harm patients and erode public trust in healthcare systems.
3Single-task baselines excluded from Table 3: all performed similar or worse than mean baseline due to low task sample size.
We are deeply grateful to Stefan Wager for his invaluable insights and extensive contributions to our discussions. We thank Emma Pierson, Kexin Huang, Kaidi Cao, Yanay Rosen, Johann Gaebler, Maria Brbic, Kefang Dong, June Vuong for helpful conversations. H.N was supported by a Stanford Knight-Hennessy Scholarship and the National Science Foundation Graduate Research Fellowship under Grant No. DGE-1656518. M.M. was supported by DARPA N660011924033 (MCS), NIH NINDS R61 NS11865, GSK, Wu Tsai Neurosciences Institute. A.S and S.O were supported by the American Slovenian Education Foundation (ASEF) fellowship. M.Y was supported by the Microsoft Research PhD fellowship. Y.R was supported by funding from GlaxoSmithKline LLC. Y.C. was supported by Stanford Graduate Fellowship and NSF IIS 2045685. We also gratefully acknowledge the support of Stanford HAI for Google Cloud Credits, DARPA under Nos. HR00112190039 (TAMI), N660011924033 (MCS); ARO under Nos. W911NF-16-1-0342 (MURI), W911NF-16-1-0171 (DURIP); NSF under Nos. OAC-1835598 (CINES), OAC-1934578 (HDR), CCF-1918940 (Expeditions), NIH under No. 3U54HG010426-04S1 (HuBMAP), Stanford Data Science Initiative, Wu Tsai Neurosciences Institute, Amazon, Docomo, GSK, Hitachi, Intel, JPMorgan Chase, Juniper Networks, KDDI, NEC, and Toshiba.
A Zero-shot Rademacher complexity and Proof of Theorem 1
A.1 Problem setup and assumptions
Let  $w \in \mathcal{W} \subseteq \mathbb{R}^e$  denote an intervention and  $x \in \mathcal{X} \subseteq \mathbb{R}^d$  denote an individual that received it. Assume the outcome to predict is a scalar  $y \in [0,1]$ . The hypothesis class is  $\mathcal{F} = \{f : (w, x) \to y\}$ . The dataset has  $n$  interventions with  $m$  independent units which received each intervention, i.e., first  $n$  i.i.d. draws from  $P_W$  and then  $m$  i.i.d. draws from  $P_X$  for each  $w^{(j)}$ . During training we have access to noisy estimate  $\tilde{y} = y + \xi$  where  $\xi$  is an independent noise with  $\mathbb{E}\xi = 0$  and  $|\xi| \leq \epsilon$  almost surely. We are tested directly on  $y$ .
The ERM is
$$
\hat {f} = \min _ {f} \hat {L} (f) = \min _ {f} \frac {1}{n m} \sum_ {j = 1} ^ {n} \sum_ {i = 1} ^ {m} (f (w ^ {(j)}, x _ {i} ^ {(j)}) - \hat {y} _ {i} ^ {(j)}) ^ {2}.
$$
The test error is
$$
L (f) = \mathbb {E} _ {w, x, y} (f (w, x) - y) ^ {2} \tag {8}
$$
and let  $f^{*} = \min_{f} L(f)$ .
We are interested in bounding the excess error  $L(\hat{f}) - L(f^{*})$ .
Our key assumption is that interventions with similar attributes  $(w)$  have similar effects in expectation. More concretely, we assume that all hypotheses in our family are smooth with respect to  $w$ :
$$
\forall f \in \mathcal {F}, \mathbb {E} _ {w, x} \left[ \left\| \frac {\partial f}{\partial w} \right\| _ {2} ^ {2} \right] \leq \beta^ {2}.
$$
Furthermore, we assume that  $P_W$  satisfies a Poincaré-type inequality:
Assumption 3. For some constant  $C$  that only depends on  $P_W$ , for any smooth function  $F$ ,
$$
V a r _ {w} [ F (w) ] \leq C \mathbb {E} \left[ \| \nabla_ {w} F (w) \| _ {2} ^ {2} \right].
$$
For example,  $P_W$  can be any of the following distributions:
- Multivariate Gaussian:  $w \in \mathbb{R}^e \sim \mathcal{N}(\mu, \Sigma)$  for some vector  $\mu \in \mathbb{R}^e$  and positive semi-definite matrix  $\Sigma \in \mathbb{R}^{e \times e}$ ;
-  $w \in \mathbb{R}^e$  has independent coordinates; each coordinate has the symmetric exponential distribution  $1 / 2e^{-|t|}$  for  $t \in \mathbb{R}$ .
-  $P_W$  is a mixture over base distributions satisfying Poincaré inequalities, and their pair-wise chi-squared distances are bounded.
-  $P_W$  is a mixture of isotropic Gaussians in  $\mathbb{R}^e$ .
-  $P_W$  is the uniform distribution over  $\mathcal{W} \subset \mathbb{R}^e$ , which is open, connected, and bounded with Lipschitz boundary.