SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Gecia Bravo-Hermsdorff* David S. Watson† Jialin Yu* Jakob Zeitler‡ Ricardo Silva*
One of the goals of causal inference is to generalize from past experiments and observational data to novel conditions. While it is in principle possible to eventually learn a mapping from a novel experimental condition to an outcome of interest, provided a sufficient variety of experiments is available in the training data, coping with a large combinatorial space of possible interventions is hard. Under a typical sparse experimental design, this mapping is ill-posed without relying on heavy regularization or prior distributions. Such assumptions may or may not be reliable, and can be hard to defend or test. In this paper, we take a close look at how to warrant a leap from past experiments to novel conditions based on minimal assumptions about the factorization of the distribution of the manipulated system, communicated in the well-understood language of factor graph models. A postulated interventional factor model (IFM) may not always be informative, but it conveniently abstracts away a need for explicitly modeling unmeasured confounding and feedback mechanisms, leading to directly testable claims. Given an IFM and datasets from a collection of experimental regimes, we derive conditions for identifiability of the expected outcomes of new regimes never observed in these training data. We implement our framework using several efficient algorithms, and apply them on a range of semi-synthetic experiments.
Causal inference is a fundamental problem in many sciences, such as clinical medicine [8, 73, 69] and molecular biology [66, 41, 29]. For example, causal inference can be used to identify the effects of chemical compounds on cell types [73] or determine the underlying mechanisms of disease [52].
One particular challenge in causal inference is generalization — the ability to extrapolate knowledge gained from past experiments and observational data to previously unseen scenarios. Consider a laboratory that has performed several gene knockouts and recorded subsequent outcomes. Do they have sufficient information to predict how the system will behave under some new combination(s) of knockouts? Conducting all possible experiments in this setting would be prohibitively expensive and time consuming. A supervised learning method could, in principle, map a vector representation of the design to outcome variables of interest. However, past experimental conditions may be too sparsely distributed in the set of all possible assignments, and such a direct supervised mapping would require leaps of faith about how assignment decisions interact with the outcome, even if under the guise of formal assumptions such as linearity.
In this paper, we propose a novel approach to the task of intervention generalization, i.e., predicting the effect of unseen treatment regimes. We rely on little more than a postulated factorization of the
*Department of Statistical Science, University College London.
†Department of Informatics, King's College London.
$\ddagger$ Department of Computer Science, University College London.
Figure 1: (a) A dynamic process where only the final step is (simultaneously) observed. Intervention variables are represented as white squares, random variables as circles, and dashed lines indicate hidden variables. (b) A bird's eye unstructured view of the sampling process, where an intervention vector dictates  $\sigma$  the distribution obtained at an agreed-upon time frame  $T$ .
distribution describing how intervention variables  $\sigma$  interact with a random process  $X$ , and a possible target outcome  $Y$ .
Related setups appear in the causal modeling literature on soft interventions [18], including its use in methods such as causal bandits [45] and causal Bayesian optimization [4] (see Section 6). However, these methods rely on directed acyclic graphs (DAGs), which may be hard to justify in many applications. Moreover, in the realistic situation where the target of a real-world intervention is a set of variables [26] and not only the "child" variable within a DAG factor, the parent-child distinction is blurred and previous identification results do not apply. Without claiming that our proposal is appropriate for every application, we suggest the following take on intervention generalization: model a causal structure as a set of soft constraints, together with their putative (arbitrary but local) modifications by external actions. Our interventional factor model (IFM) leads to provable intervention generalization via a factor graph decomposition which, when informative, can be tested without further assumptions beyond basic relations of conditional independence. IFMs are fully agnostic with regards to cycles or hidden variables, in the spirit of Dawid [23]'s decision-theoretic approach to causal inference, where the key ingredient boils down to statements of conditional independence among random and intervention variables.
Our primary contributions are as follows. (1) We introduce the *interventional factor model* (IFM), which aims to solve intervention generalization problems using only claims about which interventions interact with which observable random variables. (2) We establish necessary and sufficient conditions for the identifiability of treatment effects within the IFM framework. (3) We adapt existing results from conformal inference to our setting, providing distribution-free predictive intervals for identifiable causal effects with guaranteed finite sample coverage. (4) We implement our model using efficient algorithms, and apply them to a range of semi-synthetic experiments.
Data assumptions. Fig. 1(a) illustrates a generative process common to many applications and particularly suitable to our framework: a perturbation, here represented by a set of interventional variables  $\sigma$ , is applied and a dynamic feedback loop takes place. Often in these applications (e.g., experiments in cell biology [66] and social science [59]), the sampling of this process is highly restrictive, and the time-resolution may boil down to a single snapshot, as illustrated by Fig. 1(b). We assume that data is given as samples of a random vector  $X$  collected cross-sectionally at a well-defined time-point (not necessarily at equilibrium) under a well-defined intervention vector  $\sigma$ . The importance of establishing a clear sampling time in the context of graphical causal models is discussed by [21]. The process illustrated in Fig. 1(a) may suggest no particular Markovian structure at the time the snapshot is taken. However, in practice, it is possible to model the black-box sampling process represented in Fig. 1(b) in terms of an energy function representing soft constraints [48]. These could be the result of particular equilibrium processes [47], including deterministic differential equations [40], or empirically-verifiable approximations [59].
The role of a causal model is to describe how intervention  $\sigma$  locally changes the energy function. In the context of causal DAG models, sometimes this takes the guise of "soft interventions" and other variations that can be called "interventions on structure" [43, 55] or edge interventions [70]. Changes of structural coefficients in possibly cyclic models have also been considered [37]. The model family we propose takes this to the most abstract level, modeling energy functions via recombinations of
Figure 2: Examples of distinct causal graphical models, expressing different factorization assumptions. Random variables are represented as circles, intervention variables as white squares, and factors as black squares. (a) A directed acyclic graph (DAG) with explicit intervention variables. (b) The corresponding interventional factor model (IFM) for (a). (c) A Markov random field (MRF) with interventional variables, thus, forming a chain graph.
local interventional effects without acyclicity constraints, as such constraints should not be taken for granted [22] and are at best a crude approximation for some problems at hand (e.g., [66]).
Background and notation. The notion of an intervention variable in causal inference encodes an action that modifies the distribution of a system of random variables. This notion is sometimes brought up explicitly in graphical formulations of causal models [60, 72, 23]. To formalize it, let  $\sigma$  denote a vector of intervention variables (also known as regime indicators), with each  $\sigma_{i}$  taking values in a finite set  $\{0,1,2,\dots ,\aleph_{i} - 1\}$ . A fully specified vector of intervention variables characterizes a regime or an environment (we use these terms interchangeably). Graphically, we represent intervention variables using white squares and random variables using circles (Fig. 2). We use subscripts to index individual (random or interventional) variables and superscripts to index regimes. For instance,  $X_{i}^{j}$  denotes the random variable  $X_{i}$  under regime  $\sigma^j$ , while  $\sigma_i^j$  denotes the  $i$ th intervention variable of the  $j$ th environment. (Note that the order of the variables is arbitrary.) We occasionally require parenthetical superscripts to index samples, so  $x_{i}^{j(k)}$  denotes the  $k$ th sample of  $X_{i}$  under regime  $\sigma^j$ .
For causal DAGs, Pearl's  $do$  operator [60] denotes whether a particular random variable  $X_{i}$  is manipulated to have a given value (regardless of the values of its parents). As an example, consider a binary variable  $X_{i} \in \{True, False\}$ . We can use a (categorical) intervention variable  $\sigma_{i}$  to index the two "do-interventional" regimes:  $\sigma_{i} = 1$  denotes "do  $(X_{i} = True)$ ", and  $\sigma_{i} = 2$  denotes "do  $(X_{i} = False)$ ". We reserve  $\sigma_{i} = 0$  to denote the choice of "no manipulation", also known as the "observational" regime. However, one need not think of this  $\sigma_{i} = 0$  setting as fundamentally different than the others; indeed, it is convenient to treat all regimes as choices of data generating process. Moreover, there is no need for intervention variables to correspond to deterministic settings of random variables; they may in principle describe any well-defined change in distribution, such as stochastic or conditional interventions [23, 17]. As a stochastic example,  $\sigma_{i} = 3$  could correspond to randomly choosing  $do(X_{i} = True)$  30% of the time and  $do(X_{i} = False)$  70% of the time. As a conditional example,  $\sigma_{i} = 4$  could mean "if parents of  $X_{i}$  satisfy a given condition,  $do(X_{i} = True)$ , otherwise do nothing". When  $X_{i}$  can take more values, the options for interventional regimes become even more varied. The main point to remember is that  $\sigma_{i} = 1, \sigma_{i} = 2, \ldots$  should be treated as distinct categorical options, and we reserve  $\sigma_{i} = 0$  to denote the "observational" case of "no manipulation".
As discussed in the previous section, the probability density/mass function  $p(x; \sigma)$  can be the result of a feedback process that does not naturally fit a DAG representation. Indeed, a growing literature in causal inference carefully considers how DAGs may give rise to equilibrium distributions (e.g., [47, 12, 11]), or margins of continuous-time processes (e.g., [57]). However, they come with considerable added complexity of assumptions to ensure identifiability. In this work, we abstract away all low-level details about how an equilibrium distribution comes to be, and instead require solely a model for how a distribution  $p(x)$  factorizes as a function of  $\sigma$ . These assumptions are naturally formulated as a factor graph model [44] augmented with intervention variables, which we call an interventional factor model (IFM).
4In general, there can be infinitely many choices for interventional settings. However, for the identifiability results in the next section, we cover the finite case only, as it removes the need for smoothness assumptions on the effect of intervention levels.
Problem statement. We are given a space  $\Sigma$  of possible values for an intervention vector  $\sigma$  of dimension  $d$ , making  $|\Sigma| \leq \prod_{i=1}^{d} \aleph_{i}$ . For a range of training regimes  $\Sigma_{\mathrm{train}} \subset \Sigma$ , we are given collection of datasets  $\mathcal{D}^1, \mathcal{D}^2, \ldots, \mathcal{D}^t$ , with dataset  $\mathcal{D}^j$  collected under environment/ regime  $\sigma^j \in \Sigma_{\mathrm{train}}$ . The goal is to learn  $p(x; \sigma^{\star})$ , and  $\mu(\sigma^{\star}) := \mathbb{E}[Y; \sigma^{\star}]$ , for all test regimes  $\sigma^{\star} \in \Sigma_{\mathrm{test}} = \Sigma \backslash \Sigma_{\mathrm{train}}$ .
In each training dataset, we measure a sample of post-treatment i.i.d. draws of some  $m$ -dimensional random vector  $X$  (possibly with  $m \neq d$ , as there is no reason to always assume a one-to-one mapping between intervention and random variables), and optionally an extra outcome variable  $Y$ . The data generating process  $p(x; \sigma^j)$  is unknown. We assume  $Y \perp \sigma \mid X$  for simplicity<sup>5</sup>, which holds automatically in cases where  $Y$  is a known deterministic summary of  $X$ . We are also given a factorization of  $p(x; \sigma)$ ,
$$
p (x; \sigma) \propto \prod_ {k = 1} ^ {l} f _ {k} \left(x _ {S _ {k}}; \sigma_ {F _ {k}}\right), \forall \sigma \in \Sigma , \tag {1}
$$
where  $S_{k} \subseteq [m]$  and  $F_{k} \subseteq [d]$  define the causal structure. The (positive) functions  $f_{k}(\cdot ;\cdot)$  are unknown. The model  $p(y\mid x)$  can also be unknown, depending on the problem.
As illustrated in Figs. 2(a-b), such an intervention factor model (IFM) can also encode (a relaxation of) the structural constraints implied by DAG assumptions. In particular, note the one-to-one correspondence between the DAG factorization in Fig. 2(a) and the factors in the IFM in Fig. 2(b):  $f_{1}(x_{1};\sigma_{1}) \coloneqq p(x_{1};\sigma_{1}), f_{2}(x_{1},x_{2};\sigma_{2}) \coloneqq p(x_{2} \mid x_{1};\sigma_{2}),$  and  $f_{3}(x_{1},x_{2},x_{3};\sigma_{3}) \coloneqq p(x_{3} \mid x_{1},x_{2};\sigma_{3})$ . This explicitly represents that the conditional distribution of  $x_{1}$  given all other variables fully factorizes in  $\sigma$ : i.e.,  $p(x_1 \mid x_2,x_3;\sigma) \propto f_1(x_1;\sigma_1)f_2(x_1,x_2;\sigma_2)f_3(x_1,x_2,x_3;\sigma_3)$ . Such factorization is not enforced by usual parameterizations of Markov random fields (i.e., graphical models with only undirected edges) or chain graphs (graphical models with acyclic directed edges and undirected edges) [25] such as the example shown in Fig. 2(c).
Scope and limitations. The factorization in Eq. (1) may come from different sources, e.g., from knowledge about physical connections (it is typically the case that one is able to postulate which variables are directly or only indirectly affected by an intervention), or as the result of structure learning methods (e.g., [1]). For structure learning, faithfulness-like assumptions [72] are required, as conditional independencies discovered under configurations  $\Sigma_{\mathrm{train}}$  can only be extrapolated to  $\Sigma_{\mathrm{test}}$  by assuming that independencies observed over particular values of  $\sigma$  can be generalized across all regimes. We do not commit to any particular structure learning technique, and refer to the literature on eliciting and learning graphical structure for a variety of methods [42]. In Appendix A, we provide a general guide on structure elicitation and learning, and a primer on causal modeling and reasoning based solely on abstract conditional independence statements, without a priori commitment to a particular family of graphical models [23].
Even then, the structural knowledge expressed by the factorization in Eq. (1) may be uninformative. Depending on the nature of  $\Sigma_{\mathrm{train}}$  and  $\Sigma_{\mathrm{test}}$ , it may be the case that we cannot generalize from training to test environments, and  $p(x; \sigma^{\star})$  is unidentifiable for all  $\sigma^{\star} \in \Sigma_{\mathrm{test}}$ . However, all methods for causal inference rely on a trade-off between assumptions and informativeness. For example, unmeasured confounding may imply no independence constraints, but modeling unmeasured confounding is challenging, even more so for equilibrium data without observable dynamics. If we can get away with pure factorization constraints implied by an array of experimental conditions and domain knowledge, we should embrace this opportunity. This is what is done, for instance, in the literature on causal bandits and causal Bayesian optimization [45, 49, 4, 75], which leverage similar assumptions to decide what to do next. However, we are not proposing a method for bandits, Bayesian optimization, or active learning. The task of estimating  $p(x; \sigma^{\star})$  and  $\mu(\sigma^{\star})$  for a novel regime is relevant in itself.
3 Interventional Factor Model: Identification
We define our identification problem as follows: given the population distributions  $\mathbb{P}(\Sigma_{\mathrm{train}}) \coloneqq \{p(x; \sigma^1), p(x; \sigma^2), \ldots, p(x; \sigma^t)\}$  for the training regimes  $\sigma^j \in \Sigma_{\mathrm{train}}$ , and knowledge of the factorization assumptions as given by Eq. (1), can we identify a given  $p(x; \sigma^\star)$  corresponding to some unobserved test regime  $\sigma^\star \in \Sigma_{\mathrm{test}}$ ?
5. Otherwise, we can just define  $Y \coloneqq X_{m}$ , and use the identification results in Section 3 to check whether  $Y$  can be predicted from  $\sigma$ .
c.
Figure 3: (a) An interventional factor model (IFM) with three (binary) intervention variables. (b) The junction tree of the  $\sigma$ -graph associated with the IFM in (a): intervention variables are arranged as a hypergraph, where the hypervertices represent the sets of intervention variables that share a factor and the edge represents the overlap between the two sets of intervention variables. (c) A table displaying the assignments for regimes in  $\Sigma_{\mathrm{train}}$ . From this  $\Sigma_{\mathrm{train}}$  and the assumptions in (a), it is possible to generalize to the two missing regimes, i.e.,  $\sigma_1 = \sigma_2 = \sigma_3 = 1$ , and  $\sigma_1 = \sigma_3 = 1$ ,  $\sigma_2 = 0$ .
This identification problem is not analyzed in setups such as causal Bayesian optimization [4, 75], which build on DAGs with interventions targeting single variables. Without this analysis, it is unclear to what extent the learning might be attributed to artifacts due to the choice of prior distribution or regularization, particularly for a sparsely populated  $\Sigma_{\mathrm{train}}$ . As an extreme case that is not uncommon in practice, if we have binary intervention variables and never observe more than one  $\sigma_{i}$  set to 1 within the same experiment, it is unclear why we should expect a non-linear model to provide information about pairs of assignments  $\sigma_{i} = \sigma_{j} = 1$  using an off-the-shelf prior or regularizer. Although eventually a dense enough process of exploration will enrich the database of experiments, this process may be slow and less suitable to situations where the goal is not just to maximize some expected reward but to provide a more extensive picture of the dose-response relationships in a system. While the smoothness of  $p(x;\sigma)$  as a function of  $\sigma$  is a relevant domain-specific information, it complicates identifiability criteria without further assumptions. In what follows, we assume no smoothness conditions, meaning that for some pair  $p(x;\sigma^a), p(x;\sigma^b)$ , with  $a \neq b$ , these probability density/mass functions are allowed to be arbitrarily different. Such setting is particularly suitable for situations where interventions do take categorical levels, with limited to no magnitude information.
Identification: preliminaries. Before introducing the main result of this section, let us start by considering the toy example displayed in Fig. 3: Fig. 3(a) shows the graphical model (how  $X$  could be factorized is not relevant here), and Fig. 3(c) shows the assignments for the training regimes  $\Sigma_{\mathrm{train}}$  (all intervention variables are binary). This training set lacks the experimental assignment  $\sigma_1 = \sigma_2 = \sigma_3 = 1$ . However, this regime is implied by the model and  $\Sigma_{\mathrm{train}}$ . To see this, consider the factorization  $p(x; \sigma) \propto f_1(x; \sigma_1, \sigma_2) f_2(x; \sigma_2, \sigma_3)$ , it implies:
$$
\frac {p (x ; (1 , 1 , 0))}{p (x ; (0 , 1 , 0))} \propto \frac {f _ {1} (x ; (1 , 1))}{f _ {1} (x ; (0 , 1))},
$$
where we used  $(1,1,0)$  etc. to represent the  $\sigma$  assignments. Because both  $(1,1,0)$  and  $(0,1,0)$  are in  $\Sigma_{\mathrm{train}}$ , the ratio is identifiable up to a multiplicative constant. Moreover, multiplying and dividing the result by  $f_{2}(x;(1,1))$ , we get:
$$
\frac {p (x ; (1 , 1 , 0))}{p (x ; (0 , 1 , 0))} \propto \frac {f _ {1} (x ; (1 , 1))}{f _ {1} (x ; (0 , 1))} \frac {f _ {2} (x ; (1 , 1))}{f _ {2} (x ; (1 , 1))} \propto \frac {p (x ; (1 , 1 , 1))}{p (x ; (0 , 1 , 1))},
$$
from which, given that  $(0,1,1)$  is also in  $\Sigma_{\mathrm{train}}$ , we can derive  $p\big(x;(1,1,1)\big)$ .
3.1 Message Passing Formulation
The steps in this reasoning can be visualized in Fig. 3(b). The leftmost hypervertex represents  $(\sigma_{1},\sigma_{2})$  and suggests  $\sigma_{1}$  can be isolated from  $\sigma_{3}$ . The first ratio  $p(x;(1,1,0)) / p(x;(0,1,0))$  considers three roles:  $\sigma_{1}$  can be isolated (it is set to a "baseline" of 0 in the denominator) and  $\sigma_{3}$  is yet to be considered (it is set to the baseline value of 0 in the numerator). This ratio sets the stage for the next step where  $(\sigma_{2},\sigma_{3})$  is "absorbed" in the construction of the model evaluated at  $(1,1,1)$ . The entries in  $\Sigma_{\mathrm{train}}$  were chosen so that we see all four combinations of  $(\sigma_{1},\sigma_{2})$  and all four combinations of
As well as the assignment  $\sigma_{1} = \sigma_{3} = 1, \sigma_{2} = 0$ , which can be recovered by an analogous argument.
$(\sigma_{2}, \sigma_{3})$ , while avoiding the requirement of seeing all eight combinations of  $(\sigma_{1}, \sigma_{2}, \sigma_{3})$ . How to generalize this idea is the challenge. Next, we present our first proposed solution.
**Definitions.** Before we proceed, we need a few definitions. First, recall that we represent the baseline (or "observational") regime as  $\sigma_{1} = \sigma_{2} = \dots = \sigma_{d} = 0$ . This describes data captured under a default protocol, e.g., a transcriptomic study in which no genes are knocked down. Let  $\Sigma_{[Z]}^{0}$  denote the set of all environments  $\sigma^{j} \in \Sigma$  such that  $\sigma_{i}^{j} = 0$  if  $i \notin Z$ . For instance, for the example in Fig. 3, a set  $Z := \{2\}$  implies that  $\Sigma_{[Z]}^{0} = \{(0,0,0), (0,1,0)\}$ . Finally, let  $\sigma^{[Z(\star)]}$  be the intervention vector given by  $\sigma_{i} = \sigma_{i}^{\star}$ , if  $i \in Z$ , and 0 otherwise. In what follows, we also use of the concepts of graph decompositions, decomposable graphs and junction trees, as commonly applied to graphical models [46, 20]. In Appendix B, we review these concepts. An IFM  $\mathcal{I}$  with intervention vertices  $\sigma_{1}, \ldots, \sigma_{d}$  has an associated  $\sigma$ -graph denoted by  $\mathcal{G}_{\sigma(\mathcal{I})}$ , defined as an undirected graph with vertices  $\sigma_{1}, \ldots, \sigma_{d}$ , and where edge  $\sigma_{i} - \sigma_{j}$  is present if and only if  $\sigma_{i}$  and  $\sigma_{j}$  are simultaneously present in at least one common factor  $f_{k}$  in  $\mathcal{I}$ . For example, the  $\sigma$ -graph of the IFM represented in Fig. 3(a) is  $\sigma_{1} - \sigma_{2} - \sigma_{3}$ . This is a decomposable graph with vertex partition  $A = \{\sigma_{1}\}$ ,  $B = \{\sigma_{2}\}$  and  $C = \{\sigma_{3}\}$ . As this  $\sigma$ -graph is an undirected decomposable graph, it has a junction tree, which we depict in Fig. 3(b).