SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hamed Nilforoshan\*1 Michael Moor\*1 Yusuf Roohani2 Yining Chen1 Anja Šurina3 Michihiro Yasunaga1 Sara Oblak4 Jure Leskovec1
$^{1}$ Department of Computer Science, Stanford University
$^{2}$ Department of Biomedical Data Science, Stanford University
$^{3}$ School of Computer and Communication Sciences, EPFL
$^{4}$ Department of Computer Science, University of Ljubljana
Correspondence to: hamedn@cs.stanford.edu, mdmoor@cs.stanford.edu, jure@cs.stanford.edu
Predicting how different interventions will causally affect a specific individual is important in a variety of domains such as personalized medicine, public policy, and online marketing. There are a large number of methods to predict the effect of an existing intervention based on historical data from individuals who received it. However, in many settings it is important to predict the effects of novel interventions (e.g., a newly invented drug), which these methods do not address. Here, we consider zero-shot causal learning: predicting the personalized effects of a novel intervention. We propose CaML, a causal meta-learning framework which formulates the personalized prediction of each intervention's effect as a task. CaML trains a single meta-model across thousands of tasks, each constructed by sampling an intervention, its recipients, and its nonrecipients. By leveraging both intervention information (e.g., a drug's attributes) and individual features (e.g., a patient's history), CaML is able to predict the personalized effects of novel interventions that do not exist at the time of training. Experimental results on real world datasets in large-scale medical claims and cell-line perturbations demonstrate the effectiveness of our approach. Most strikingly, CaML's zero-shot predictions outperform even strong baselines trained directly on data from the test interventions.
Personalized predictions about how an intervention will causally affect a specific individual are important across many high impact applications in the physical, life, and social sciences. For instance, consider a doctor deciding whether or not to prescribe a drug to a patient. Depending on the patient, the same drug could either (a) cure the disease, (b) have no effect, or (c) elicit a life-threatening adverse reaction. Predicting which effect the drug will have for each patient could revolutionize healthcare by enabling personalized treatments for each patient.
The causal inference literature formalizes this problem as conditional average treatment effects (CATE) estimation, in which the goal is to predict the effect of an intervention, conditioned on patient characteristics  $(X)$ . When natural experiment data is available, consisting of individuals who already did and did not receive an intervention, a variety of CATE estimators exist to accomplish this task [1, 3, 16, 23, 30, 34, 44, 55, 36, 70]. These methods can then predict the effect of an existing intervention  $(W)$  on a new individual  $(X^{\prime})$ .
However, in many real-world applications natural experiment data is entirely unavailable, and yet CATE estimation is critical. For instance, when new drugs are discovered, or new government policies
*Equal contribution. Code is available at: https://github.com/snap-stanford/caml/
are passed, it is important to know the effect of these novel interventions on individuals and subgroups in advance, i.e., before anybody is treated. There is thus a need for methods that can predict the effect of a novel intervention  $(W^{\prime})$  on a new individual  $(X^{\prime})$  in a zero-shot fashion, i.e., without relying on any historical data from individuals who received the intervention.
Generalizing to novel interventions is especially challenging because it requires generalizing across two dimensions simultaneously: to new interventions and new individuals. This entails efficiently "aligning" newly observed interventions to the ones previously observed in the training data.
Present work. Here, we first formulate the zero-shot CATE estimation problem. We then propose CaML (Causal Meta-learning), a general framework for training a single meta-model to estimate CATE across many interventions, including novel interventions that did not exist at the time of model training (Figure 1). Our key insight is to frame CATE estimation for each intervention as a separate meta-learning task. For each task observed during training, we sample a retrospective natural experiment consisting of both (a) individuals who did receive the intervention, and (b) individuals who did not receive the intervention. This natural experiment data is used to estimate the effect of the intervention for each individual (using any off-the-shelf CATE estimator), which serves as the training target for the task.
In order to achieve zero-shot generalization to new interventions, we include information  $(W)$  about the intervention (e.g., a drug's attributes), in the task. We then train a single meta-model which fuses intervention information with individual-level features  $(X)$  to predict the intervention's effect. Our approach allows us to predict the causal effect of novel interventions, i.e., interventions without sample-level training data, such as a newly discovered drug (Figure 1). We refer to this capability as zero-shot causal learning.
In our experiments, we evaluate our method on two real-world datasets—breaking convention with the CATE methods literature which typically relies on synthetic and semi-synthetic datasets. Our experiments show that CaML is both scalable and effective, including the application to a large-scale medical dataset featuring tens of millions of patients. Most strikingly, CaML's zero-shot performance exceeds even strong baselines that were trained directly on data from the test interventions. We further discover that CaML is capable of zero-shot generalization even under challenging conditions: when trained only on single interventions, at inference time it can accurately predict the effect of combinations of novel interventions. Finally, we explain these findings, by proving a zero-shot generalization bound.
We discuss recent work which is most closely related to zero-shot causal learning, and provide an extended discussion of other related work in Appendix B. Most CATE estimators do not address novel interventions, requiring that all considered interventions be observed during training. A notable exception is recent methods which estimate CATE for an intervention using structured information about its attributes [25, 35]. In principle, these methods can also be used for zero-shot predictions. These methods estimate CATE directly from the raw triplets  $(W, X, Y)$ , without considering natural experiments, by tailoring specific existing CATE estimators (the S-learner [44] and Robinson decomposition [55], respectively) to structured treatments. The main drawback of these approaches is that they are inflexible, i.e., they are restricted to using a single estimator and are unable to take advantage of the recent advances in the broader CATE estimation literature (e.g., recently developed binary treatment estimators [16, 21, 40]). This is a limitation because any single CATE estimator can be unstable across different settings [15]. Notably, the estimators which these methods build on have already been shown to result in high bias in many domains [44, 37, 10, 16]. Likewise, we find that these methods struggle with zero-shot predictions (Section 6). CaML's key difference from prior work is that we construct a separate task for each training intervention by synthesizing natural experiments. This allows us to (a) flexibly wrap any existing CATE estimator to obtain labels for each task, and thus take advantage of the most recent CATE estimation methods and (b) leverage meta-learning, which requires task-structured data. Consequently, CaML is able to achieve strong zero-shot performance (Section 6).
3 Background: single-intervention CATE estimation
Each task in the CaML framework consists of estimating conditional average treatment effects (CATEs) for a single binary treatment. In this section, we first provide background on CATE
Figure 1: Overview of the zero-shot causal learning problem. Each individual has features  $(X)$ , an intervention with features  $(W)$ , and an outcome  $(Y)$ . Lightning bolts  $(\mathcal{F})$  represent interventions (e.g. drugs). The personalized effect of an intervention  $(\tau)$  is always unobserved. The goal is to predict the  $\tau$  for a novel intervention  $(W^{\prime})$  and individual  $(X^{\prime})$  that did not exist during training.
estimation under this simple case of a single treatment  $(W)$  and outcome  $(Y)$ , and subsequently generalize it to our zero-shot setting. Under a single intervention and outcome, we consider  $n$  independent observations  $P_{1},\ldots ,P_{n}$  drawn from a distribution  $\mathcal{P}$ . For unit  $i = 1,\dots,n$ ,  $P_{i} = (W_{i},X_{i},Y_{i})\sim \mathcal{P}$  collects: a binary or continuous outcome of interest  $Y_{i}\in \mathcal{V}\subset \mathbb{R}$ , instance features (i.e., pre-treatment covariates)  $X_{i}\in \mathcal{X}\subset \mathbb{R}^{d}$ , and a treatment-assignment indicator  $W_{i}\in \{0,1\}$ . We use the Neyman-Rubin potential outcomes framework [33], in which  $Y_{i}(1),Y_{i}(0)$  reflect the outcome of interest either under treatment  $(W_{i} = 1)$ , or under control  $(W_{i} = 0)$ , respectively. In our running medical example,  $Y_{i}(1)$  is the health status if exposed to the drug, and  $Y_{i}(0)$  is the health status if not exposed to the drug. Notably, the fundamental problem of causal inference is that we only observe one of the two potential outcomes, as  $Y_{i} = W_{i}\cdot Y_{i}(1) + (1 - W_{i})\cdot Y_{i}(0)$  (e.g., either health status with or without drug exposure can be observed for a specific individual, depending on whether they are prescribed the drug). However, it is possible to make personalized decisions by estimating treatment effects that are tailored to the attributes of individuals (based on features  $X$ ). Thus, we focus on estimating  $\tau (x)$ , known as the conditional average treatment effect (CATE):
$$
\operatorname {C A T E} = \tau (x) = \mathbb {E} _ {\mathcal {P}} \left[ Y (1) - Y (0) \mid X = x \right] \tag {1}
$$
A variety of methods have been developed to estimate  $\tau(x)$  from observational data [16]. These rely on standard assumptions of unconfoundedness, consistency, and overlap [52]. Unconfoundedness: there are no unobserved confounders, i.e.  $Y_{i}(0), Y_{i}(1) \perp W_{i} \mid X_{i}$ . Consistency:  $Y_{i} = Y_{i}(W_{i})$ , i.e. treatment assignment determines whether  $Y_{i}(1)$  or  $Y_{i}(0)$  is observed. Overlap: Treatment assignment is nondeterministic, such that for all  $x$  in support of  $X$ :  $0 < P(W_{i} = 1 \mid X_{i} = x) < 1$ .
4 Zero-shot causal learning
In many real-world settings (e.g. drugs, online A/B tests) novel interventions are frequently introduced, for which no natural experiment data are available. These settings require zero-shot CATE estimates. The zero-shot CATE estimation problem extends the prior section, except the intervention variable  $W_{i}$  is no longer binary, but rather contains rich information about the intervention:  $W_{i}\in \mathcal{W}\subset \mathbb{R}^{e}$  (e.g., a drug's chemistry), where  $W_{i} = 0$  corresponds to a sample that did not receive any intervention. Thus, each intervention value  $w$  has its own CATE function that we seek to estimate:
$$
\mathrm {C A T E} _ {w} = \tau_ {w} (x) = \mathbb {E} _ {\mathcal {P}} \left[ Y (w) - Y (0) \mid X = x \right], \tag {2}
$$
During training, we observe  $n$  independent observations  $P_{1},\ldots ,P_{n}$  drawn from a distribution  $\mathcal{P}$ . Each  $P_{i} = (W_{i},X_{i},Y_{i})\sim \mathcal{P}$ . Let  $\mathcal{W}_{seen}$  be set of all interventions observed during training. The zero-shot CATE estimation task consists of estimating CATE for a novel intervention that was never observed during training:
Problem 1 (Zero-shot CATE estimation). Given  $n$  training observations  $(W_{1},X_{1},Y_{1}),\ldots ,(W_{n},X_{n},Y_{n})$  drawn from  $\mathcal{P}$  containing intervention information, individual features, and outcomes... estimate  $\tau_{w^{\prime}}(x)$  for a novel intervention  $w^{\prime}\notin \mathcal{W}_{seen}$
This problem formulation extends in a straightforward manner to combinations of interventions, by allowing a single intervention  $W_{i}$  to consist of a set of intervention vectors. CaML supports combinations of interventions, as we elaborate on in Section 4.1
Figure 2: Visual illustration of the CaML (causal meta-learning) framework. (1) We sample a task (i.e., an intervention) and a natural experiment from the training data consisting of individuals who either received the intervention  $(\mathrm{W} = \{\mathbf{\theta}\})$ , or did not  $(\mathrm{W} = \{\cdot\})$ . Each individual has features  $(X)$  and an outcome  $(Y)$ , and the intervention also has information  $(W)$  (e.g., a drug's attributes). (2) For each individual we estimate the effect of the intervention on the outcome (pseudo-outcomes  $\tilde{\tau}$ ). (3) We predict an individual's pseudo-outcomes  $\tilde{\tau}$  using a model that fuses  $X$  and  $W$ . CaML is trained by repeating this procedure across many tasks and corresponding natural experiments.
CaML overview. We propose a novel framework for estimating CATE across multiple interventions, even including ones that were never encountered during training. Our framework consists of three key components (Figure 2). First, we formulate CATE estimation as a meta-learning problem in which each task corresponds to the CATE estimation for a unique intervention. A task dataset for a given intervention is constructed by sampling a natural experiment of all individuals who received the intervention, and a sample of individuals who did not. Tasks are augmented with intervention information  $(W)$ . Synthesizing these natural experiments allows us to compute a noisy CATE label  $\tilde{\tau}$  using any off-the-shelf estimator ( $\tilde{\tau}$  is referred to as pseudo-outcomes by the causal inference literature [16]). Finally, we train a single meta-model to predict these labels using individual-level  $(X)$  and intervention-level  $(W)$  information, such that it is able to generalize to novel tasks, i.e., estimating CATE for novel interventions.
The CaML framework incorporates three important design considerations: (1) Single meta-model. In domains such as electronic health records and online marketing, we observe that large-scale datasets contain thousands of interventions with rich feature information  $(W)$ . Instead of training a separate model for each intervention, CaML trains a single meta-model that can estimate CATE across all interventions. This approach lets us leverage shared structure across tasks and generalize to novel interventions that were not present during training. (2) Pseudo-outcomes. Instead of directly modeling the response surfaces  $\mathbb{E}[Y(w) \mid X = x]$  and  $\mathbb{E}[Y(0) \mid X = x]$ , we use pseudo-outcomes for each intervention to train our model. This approach is informed by recent studies indicating bias in estimating CATE from direct predictions of observed outcomes [10, 44]. CaML outperforms strong baselines that meta-learn  $Y(w)$  and  $Y(0)$  directly, as demonstrated in our experiments (see Tables 2 and 3, rows S-learner and T-learner with meta-learning). (3) Discrete tasks from continuous interventions. CaML takes advantage of the extensive literature on CATE estimation for single, binary interventions. By creating a natural experiment for each intervention, CaML taps into this literature and benefits from the high performance of recently developed nonparametric CATE estimators [16, 55, 44].
CaML identifies CATE for novel interventions under the assumptions that: (1) for each observed intervention  $w$ ,  $\tau_w(x)$  is identifiable under the binary treatment assumptions (unconfoundedness, consistency, and overlap) in Section 3. This allows for valid training labels for each task. (2)  $\tau_w(x) = \tau(w, x)$ , i.e., a global function  $\tau(w, x)$  unifies all intervention-specific CATE functions, (3)  $\tau(w, x)$  is continuous in  $w$ . This allows the model to smoothly extrapolate the treatment effect to new interventions that are close to observed interventions in the intervention space. Lastly, (4)  $W$  follows a continuous distribution.
4.1 Meta-dataset
We formulate CATE estimation as a meta-learning problem. For this, each task refers to CATE estimation for a distinct intervention. Interventions as well as tasks in our meta-dataset are jointly indexed by  $j \in \mathbb{N}$  with  $1 \leq j \leq K$ , such that we can refer to the  $j$ -th intervention information with  $w^{(j)}$ .
1. Sample natural experiment
2. Estimate pseudo-outcomes  $(\tilde{\mathfrak{t}})$
3. Predict  $\widetilde{\tau}$  from individual and intervention information
We then construct a meta-dataset  $D$  in the following way:
$$
D = \left\{\left(D _ {\text {t r e a t e d}} ^ {(j)} \cup D _ {\text {c o n t r o l}} ^ {(j)}, w ^ {(j)}\right) \right\} _ {j = 1} ^ {K}, \text {w i t h} \tag {3}
$$
$$
D _ {\text {t r e a t e d}} ^ {(j)} = \left\{\left(X _ {i}, Y _ {i}\right) \mid W _ {i} = w ^ {(j)} \right\} \text {a n d} D _ {\text {c o n t r o l}} ^ {(j)} = \left\{\left(X _ {i}, Y _ {i}\right) \mid W _ {i} = 0\right) \}. \tag {4}
$$
$D^{(j)}$  denotes the natural experiment dataset for task  $j$ , composed of a treated group (instances which received the intervention, i.e.  $W_{i} = w^{(j)}$ ) and control group (instances which did not receive any intervention, i.e.  $W_{i} = 0$ ). Each sample  $i$  represents an individual, for which the quantities  $(X_{i}, Y_{i})$  are collected as introduced in Section 3. In practice, we down-sample both groups (i.e. to 1 million samples for the treated and control groups) in our large-scale experiments.
We augment each task dataset  $D^{(j)}$  with intervention information,  $w^{(j)} \in \mathbb{R}^e$ , for zero-shot generalization to new interventions [35, 18, 87, 39]. The form of  $w^{(j)}$  varies with the problem domain — for text interventions, it could be a language model's text embedding [79, 84, 58], while biomedical treatments can be represented as nodes in a knowledge graph [8, 49]. Additionally, domain-specific features, like treatment categories from an ontology, may be included in  $w^{(j)}$ . To handle combinations of interventions (e.g., pairs of drugs), we aggregate the  $w$  for each intervention using an order-invariant pooling operation (we used the sum operator), and sample a separate natural experiment for individuals who received the full combination.
4.2 Estimating pseudo-outcomes
We next estimate the training targets for each task (i.e. intervention) in the meta-dataset. The training target  $(\tilde{\tau}^{(j)})$  is an unbiased, but noisy, estimate of CATE. More formally, for each task  $j$  (which points to the natural experiment dataset for intervention  $w^{(j)}$ ), we estimate  $\tilde{\tau}^{(j)}$ , where  $\mathbb{E}_{\mathcal{P}}[\tilde{\tau}^{(j)}|X = x] = \tau_{w^{(j)}}(x)$ . Thus,  $\tilde{\tau}_i^{(j)}$  denotes the target for the  $i$ -th sample in the  $j$ -th task (indexing will be omitted when it is clear from context). We refer to these targets as pseudo-outcomes, following prior literature [16]. For prior work on pseudo-outcomes, refer to Appendix B. In Appendix E we demonstrate why these pseudo-outcomes provide an unbiased training objective. For a detailed explanation on the necessity of using pseudo-outcomes instead of directly modeling  $Y(w)$  and  $Y(0)$ , please see [44, 16, 10].
CaML is agnostic to the specific choice of pseudo-outcome estimator. Thus, we assume a function  $\eta(D^{(j)})$  which takes as input a task dataset  $D^{(j)} \in D$  and returns a vector containing the pseudo-outcomes  $\tilde{\tau}$  for each sample in the task. We extend each task dataset  $D^{(j)}$  with the pseudo-outcomes, such that a sample holds the elements  $(X_i, Y_i, \tilde{\tau}_i)$ . Our key insight is that by collecting these pseudo-outcomes across multiple tasks, and predicting them using a combination of intervention and individual information  $(W, X)$  we can develop a CATE estimator which generalizes to novel interventions. In practice, we use the RA-learner [17] and treat pseudo-outcome estimation as a data pre-processing step (Appendix C.6).
4.3 Meta-model training
Given  $m$  target outcomes  $Y_{1},\ldots ,Y_{m}$  (e.g., different drug side effects), our goal is then to learn a model  $\Psi_{\theta}:\mathbb{R}^{e}\times \mathbb{R}^{d}\to \mathbb{R}^{m}$  that for parameters  $\theta$  minimizes
$$
\theta^ {*} = \underset {\theta} {\operatorname {a r g m i n}} \mathbb {E} _ {j \sim U (D)} \mathbb {E} _ {W, X, \tilde {\tau} \sim D ^ {(j)}} \left[ L \left(\Psi_ {\theta}\right) \right], \tag {5}
$$
where  $U(D)$  denotes the discrete uniform distribution over the tasks of the meta-dataset  $D$ , and where  $L(f)$  refers to a standard loss function between the pseudo-outcomes and the model output, i.e.,  $L(f) = (\tilde{\tau} - f(w, x))^2$ . To assess whether the model generalizes to novel tasks, we partition our meta-dataset by task, into non-overlapping subsets  $D = D_{\mathrm{train}} \cup D_{\mathrm{val}} \cup D_{\mathrm{test}}$ . During training,  $\Psi_{\theta}$  is optimized on training tasks  $D_{\mathrm{train}}$ . We validate and test this model on  $D_{\mathrm{val}}$  and  $D_{\mathrm{test}}$ , which are thus unseen during training tasks. While the CaML framework is agnostic to a specific training strategy, we based our approach (Algorithm 1) on the Reptile meta-learning algorithm [53] which we find performs better compared to straightforward empirical risk minimization (c.f. Section 6). For this, the objective is slightly modified to
$$
\theta^ {*} = \underset {\theta} {\operatorname {a r g m i n}} \mathbb {E} _ {j \sim U (D)} \left[ L \left(A _ {D ^ {j}} ^ {k} (\Psi_ {\theta})\right) \right], \tag {6}
$$
Algorithm 1 The CaML algorithm
Require: meta-dataset  $D$ , meta-model  $\Psi_{\theta}$  with initialized parameters  $\theta$ , hyperparameter  $k$ .  
for iteration  $= 1,2,\ldots,L$  do  
 $j \leftarrow \text{SAMPLETASK()}$ $D_{\text{treat}}^{(j)}, D_{\text{ctrl}}^{(j)}, w^{(j)} \leftarrow \text{QUERYTASKDATA}(j)$ $\tilde{\tau}^{(j)} \leftarrow \text{ESTIMATEPSEUDOOUTCOMES}(D_{\text{treat}}^{(j)}, D_{\text{ctrl}}^{(j)})$ $\theta' \leftarrow \text{ADAPT}((D_{\text{treat}}^{(j)}, D_{\text{ctrl}}^{(j)}), \tilde{\tau}^{(j)}, w^{(j)}, \Psi_{\theta}, k)$ $g \leftarrow \theta - \theta'$  {Reptile gradient}  
 $\theta \leftarrow \theta - \beta g$  {Gradient step for meta-model  $\Psi_{\theta}$ }  
end for  
return  $\Psi_{\theta}$   
end function
where  $A_D^k \colon \mathcal{F} \to \mathcal{F}$  represents the operator that updates a model  $f \in \mathcal{F}$  using data sampled from the dataset  $D$  for  $k$  gradient steps. This operator is defined in more detail as the ADAPT routine in Algorithm 1. Note that depending on the choice of CATE estimator, this routine iterates only over treated samples of a task dataset  $D^{(j)}$  (as in our experiments), or over all samples, including untreated ones.
4.4 CaML architecture
To parameterize  $\Psi_{\theta}$ , we propose a simple but effective model architecture (see Section 6):
$$
\Psi_ {\theta} (w, x) = \operatorname {M L P} _ {1} ([ \tilde {w}; \tilde {x} ]), \text {w i t h} \tilde {x} = \operatorname {M L P} _ {2} (x) \text {a n d} \tilde {w} = \operatorname {M L P} _ {3} (w), \tag {7}
$$
where  $[\cdot ;\cdot ]$  denotes concatenation. Equation 7 shows that the intervention information  $w$  and individual features  $x$  are encoded separately into dense vectors  $\tilde{w}$  and  $\tilde{x}$ , respectively. Our MLPs consist of layers of the form  $g(z) = z + \mathrm{ReLU}(\mathrm{Linear}(z))$ .
We now consider zero-shot causal learning from a theoretical perspective. Under simplified assumptions, we bound the prediction error in the zero-shot setting.