SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Trust Your  $\nabla$ : Gradient-based Intervention Targeting for Causal Discovery
Mateusz Olko $^{1,2,*}$  Michal Zajac $^{3,*}$  Aleksandra Nowak $^{2,3,4,*}$
Nino Scherrer5 Yashas Annadi Stefan Bauer6
Łukasz Kuciński $^{2,7}$  Piotr Miłos $^{2,8,7}$
1Warsaw University, 2IDEAS NCBR,
$^{3}$ Jagiellonian Univeristy, Faculty of Mathematics and Computer Science,
$^{4}$ Jagiellonian University, Doctoral School of Exact and Natural Sciences,
$^{5}$ ETH Zurich,  $^{6}$ Helmholtz, TU Munich,  $^{7}$ deepsense.ai,
$^{8}$ Institute of Mathematics, Polish Academy of Sciences
Inferring causal structure from data is a challenging task of fundamental importance in science. Often, observational data alone is not enough to uniquely identify a system's causal structure. The use of interventional data can address this issue, however, acquiring these samples typically demands a considerable investment of time and physical or financial resources. In this work, we are concerned with the acquisition of interventional data in a targeted manner to minimize the number of required experiments. We propose a novel Gradient-based Intervention Targeting method, abbreviated GIT, that 'trusts' the gradient estimator of a gradient-based causal discovery framework to provide signals for the intervention targeting function. We provide extensive experiments in simulated and real-world datasets and demonstrate that GIT performs on par with competitive baselines, surpassing them in the low-data regime.
Estimating causal structure from data, commonly known as causal discovery, is central to the progress of science [Pearl, 2009]. Real-world systems can often be explained as a composition of smaller parts connected by causal relationships. Understanding this underlying structure is essential for making accurate predictions about the system's behavior after a perturbation or treatment has been applied [Peters et al., 2016]. Causal discovery methods have been successfully deployed in various fields, such as biology [Sachs et al., 2005, Triantafillou et al., 2017, Glymour et al., 2019], medicine [Shen et al., 2020, Castro et al., 2020, Wu et al., 2022], earth system science [Ebert-Uphoff and Deng, 2012], or neuroscience [Sanchez-Romero et al., 2019]. In machine learning, causal decomposition has been shown to enable sample-efficient learning and fast adaptation to distribution shifts by only updating a subset of parameters [Bengio et al., 2020, Scherrer et al., 2022].
Observational data, that is the data obtained directly from the unperturbed system, are, in general, insufficient to identify a system's causal structure and only allow to determine the structure up to the so-called Markov Equivalence Class [Spirtes et al., 2000a, Peters et al., 2017]. To overcome this limited identifiability problem, causal discovery algorithms commonly leverage interventional data [Hauser and Buhlmann, 2012, Brouillard et al., 2020, Ke et al., 2019], which are acquired by gathering data from an experiment perturbing a part of the system [Spirtes et al., 2000b, Pearl, 2009]. The field of experimental design [Lindley, 1956, Murphy, 2001, Tong and Koller, 2001] is concerned with the
*These authors contributed equally. Corresponding author: mateusz.olko@gmail.com
acquisition of interventional data in a targeted manner to minimize the number of required experiments, which often requires spending a significant amount of time and physical or financial resources.
In this work, we introduce a simple and effective experimental design algorithm called Gradient-based Intervention Targeting, or GIT for short, see Figure 1. GIT can be easily combined with various gradient-based causal discovery frameworks to provide an efficient active selection of intervention targets. Our method, which is grounded in the ideas from active and curriculum learning [Settles et al., 2007, Graves et al., 2017, Ash et al., 2020], collects interventional data that induce the biggest gradient on parameters of causal structure. GIT leverages the gradient-based nature of the underlying causal discovery framework and achieves better performance than the contemporary baselines.
- We introduce GIT, which is to our knowledge, the first gradient-based intervention targeting method. Due to it's plug-and-play nature, our method can be easily combined with various gradient-based causal discovery frameworks.
- Our extensive experiments on synthetic and real-world graphs demonstrate that GIT effectively reduces the amount of interventional data needed to discover the causal structure, and performs well in the low-data regime. This makes GIT a compelling option when access to interventional data is limited.
- We provide a theoretical justification of GIT and a suite of analyses introspecting its behavior and performance.
Figure 1: Overview of GIT's usage in a gradient-based causal discovery framework. The framework infers a posterior distribution over graphs from observational and interventional data (denoted as  $\mathcal{D}_{obs}$  and  $\mathcal{D}_{int}$ ) through gradient-based optimization. The distribution over graphs and the gradient estimator  $\nabla \mathcal{L}(\cdot)$  are then used by GIT in order to score the intervention targets based on the magnitude of the estimated gradients. The intervention target with the highest score is then selected, upon which the intervention is performed. New interventional data  $\mathcal{D}_{int}^{new}$  are then collected and the procedure is repeated.
Experimental Design / Intervention Design. There are two major classes of methods for selecting optimal interventions for causal discovery. One class of approaches is based on graph-theoretical properties. Typically, a completed partially directed acyclic graph (CPDAG), describing an equivalence class of DAGs, is first specified. Then, either substructures, such as cliques or trees, are investigated and used to inform decisions [He and Geng, 2008, Eberhardt, 2008, Squires et al., 2020, Greenewald et al., 2019], or edges of a proposed graph are iteratively refined until reaching a prescribed budget [Ghassami et al., 2018, 2019, Kocaoglu et al., 2017, Lindgren et al., 2018]. One limitation of graph-theoretical approaches is that misspecification of the CPDAG at the beginning of the process can deteriorate the final solution. Another class of methods is based on Bayesian Optimal Experiment Design [Lindley, 1956], which aims to select interventions with the highest mutual information (MI) between the observations and model parameters. MI is approximated in different ways: AIT [Scherrer et al., 2021] uses F-score inspired metric to implicitly approximate MI; CBED [Tugas et al., 2022] incorporates BALD-like estimator [Houlsby et al., 2011]; ABCD [Agrawal et al., 2019] uses estimator based on weighted importance sampling. Although theoretically principled, computing mutual information suffers from approximation errors and model mismatches. Therefore, in this work, we explore using scores based on different principles.
Gradient-based Causal Structure Learning. The appealing properties of neural networks have sparked a flurry of gradient-based causal structure learning methods. The most prevalent approaches are self-supervised formulations that optimize a data-dependent scoring metric (for instance, penalized log-likelihood) to find the best causal graph  $G$ . Existing self-supervised methods that are capable (or can be extended) of incorporating interventional data can be categorized based on
the underlying optimization formulation into: (i) frameworks with a joint optimization objective [Brouillard et al., 2020, Lorch et al., 2021, Cundy et al., 2021, Annadani et al., 2021, Geffner et al., 2022, Deleu et al., 2022] and (ii) frameworks with alternating phases of optimization [Bengio et al., 2020, Ke et al., 2019, Lippe et al., 2022]. While structural and functional parameters are optimized under a joint objective in the former, the latter splits the optimization into two phases with separate objectives. All the aforementioned methods allow evaluation of gradient with respect to the structural and functional parameters with a batch of (real or hypothesized) interventional samples and can serve as a base framework for our proposed gradient-based intervention acquisition strategy.
Gradients in Active and Curriculum Learning. Gradients have been successfully used as a criterion to select data to process in previous work. Settles et al. [2007] introduce Expected Gradient Length (EGL), computed under the current belief, as a criterion for active learning. A batch active learning method introduced in Ash et al. [2020] also targets data points with high gradient magnitude, including uncertainty and diversity in the decision. In the area of curriculum learning, Graves et al. [2017] considers Gradient Prediction Gain (GPG), which is defined as the gradient's magnitude and is meant to be a proxy for expected learning progress. We take inspiration from those approaches to propose a novel usage of the gradient criterion in the field of causal discovery.
3.1 Structural Causal Models and Causal Structure Discovery
Causal relationships can be formalized using structural causal models (SCM) [Peters et al., 2017]. Each of the endogenous variables  $X = (X_{1},\ldots ,X_{n})$  is expressed as a function  $X_{i} = f_{i}(PA_{i},U_{i})$  of its direct causes  $PA_{i}\subseteq X$  and an external independent noise  $U_{i}$ . It is assumed that the assignments are acyclic and thus associated with a directed acyclic graph  $G = (V,E)$ . The nodes  $V = \{1,\dots ,n\}$  represent the random variables and the edges correspond to the direct causes, that is  $(i,j)\in E$  if and only if  $X_{i}\in PA_{j}$ . The joint distribution factorizes according to
$$
\mathbb {P} \left(X _ {1}, \dots , X _ {n}\right) = \prod_ {i = 1} ^ {n} \mathbb {P} \left(X _ {i} \mid P A _ {i}\right). \tag {1}
$$
Causal structure discovery aims to recover the ground truth graph  $G$ . The solution to this problem is not uniquely defined when having access only to observational data from the ground truth distribution  $\mathbb{P}$ . Formally, it can be determined solely up to a Markov Equivalence Class (MEC) [Spirtes et al., 2000b, Peters et al., 2017] without additional restrictive assumptions. To achieve identifiability, data from additional experiments, called interventions, need to be gathered.
A single-node intervention on  $X_{i}$  replaces the conditional distribution  $\mathbb{P}(X_i|PA_i)$  with a new distribution denoted as  $\mathbb{P}(X_i|PA_i)$ , yielding a so-called interventional distribution:
$$
\mathbb {P} _ {i} (X) \triangleq \widetilde {\mathbb {P}} \left(X _ {i} \mid P A _ {i}\right) \prod_ {j \neq i} \mathbb {P} \left(X _ {j} \mid P A _ {j}\right). \tag {2}
$$
The node  $i \in V$  is called the intervention target. An intervention that removes the dependency of a variable  $X_{i}$  on its parents, yielding  $\mathbb{P}(X_i|PA_i) = \widetilde{\mathbb{P}} (X_i)$ , is called hard. In this paper, we use data gathered by performing single-node interventions.
3.2 Online Causal Discovery and Targeting Methods
In this work, we consider an online causal discovery procedure outlined in Algorithm 1. Given a causal discovery Algorithm  $\mathcal{A}$ , the graph model  $\varphi_0$  is fitted using observational data  $\mathcal{D}_{obs}$ . Following that, batches of interventional samples are acquired iteratively and are used by the algorithm to improve the belief about the causal structure (line 7). Intervention targets are chosen by intervention targeting method  $\mathcal{M}$  to optimize the overall performance, taking into account the current belief about the graph structure encoded in  $\varphi_{i-1}$ . Below we discuss two popular choices for the method  $\mathcal{M}$  (with more details deferred to Appendix D).
Algorithm 1 ONLINE CAUSAL DISCOVERY
input causal discovery algorithm  $\mathcal{A}$  (e.g., ENCO, see Sec 4.1), intervention targeting method  $\mathcal{M}$ , number of data acquisition rounds  $T$ , observational dataset  $\mathcal{D}_{obs}$
output final parameters of graph model:  $\varphi_T$
1:  $\mathcal{D}_{int}\gets \emptyset$
2: Fit graph model  $\varphi_0$  with algorithm  $\mathcal{A}$  on  $\mathcal{D}_{obs}$
3: for round  $i = 1,2,\ldots ,T$  do
4:  $I\gets$  generate intervention targets using  $\mathcal{M}$
5:  $D_{int}^{I}\gets$  query for data from interventions  $I$
6:  $\mathcal{D}_{int} \gets \mathcal{D}_{int} \cup D_{int}^{I}$
7: Fit  $\varphi_{i}$  with algorithm  $\mathcal{A}$  on  $\mathcal{D}_{int}$  and  $\mathcal{D}_{obs}$
8: end for
Active Intervention Targeting (AIT) AIT selects the intervention target according to an  $F$ -test inspired criterion [Scherrer et al., 2021]. It assumes that the causal discovery algorithm  $\mathcal{A}$  maintains a posterior distribution over graphs (by design or using bootstrapping). To select an intervention target, a set of graphs is sampled from the posterior distribution, and interventional sample distributions are generated by intervening on each of the sampled graphs. Each potential intervention target is assigned a score by measuring the discrepancy across the corresponding interventional sample distributions.
CBED targeting Another approach to
causal discovery is approximating the posterior distribution over the possible causal DAGs. This allows using the framework of Bayesian Optimal Experimental Design to select the most informative intervention (experiment). The score of a new experiment is given by the mutual information (MI) between the interventional data due to the experiment and the current belief about the graph structure. Hence, such an approach requires estimating MI. For instance, Causal Bayesian Experimental Design (CBED) [Tigas et al., 2022] uses a BALD-like estimator [Houlsby et al., 2011] to sample batches of interventional targets.
4 GIT method
In this work, we present a new intervention targeting method GIT. GIT chooses intervention targets that induce the largest update of the parameters modeling the causal structure. Inspired by hallucinated gradients exploited by [Ash et al., 2020] we calculate gradients on imaginary data generated by the causal model, to score possible interventions for real data acquisition.
To formally introduce our method, we first describe the requirements that need to be fulfilled by a causal algorithm  $\mathcal{A}$  in order to use it with GIT. We then explain how GIT works and follow up with a discussion about causal assumptions and theoretical justification of our approach. Finally, in Section 4.1, we present a practical implementation of our method with a causal discovery algorithm  $\mathcal{A}$ , using a popular ENCO algorithm as an example.
Requirements for causal discovery algorithm  $\mathcal{A}$ . The intervention targeting method GIT can be coupled with any gradient-based causal discovery algorithm  $\mathcal{A}$  (see Algorithm 1) that fulfills the following conditions:
1.  $\mathcal{A}$  models a distribution over the causal DAGs, denoted by a family of probability measures  $\mathbb{P}_{\rho}(G)$  parameterized by  $\rho$ , that allows sampling.
2. For each causal graph  $G$ ,  $\mathcal{A}$  maintains a corresponding family of conditional distributions,  $\mathbb{P}_{G,\phi}(X_i|PA_{(i,G)})$ , parametrized by  $\phi$ , which induces the joint distribution  $\mathbb{P}_{G,\phi}$ :
$$
\mathbb {P} _ {G, \phi} (X) \triangleq \prod_ {i} \mathbb {P} _ {G, \phi} \left(X _ {i} \mid P A _ {(i, G)}\right). \tag {3}
$$
If  $G$  corresponds to the ground truth graph,  $\mathbb{P}_{G,\phi}$  approximates the ground truth distribution over  $X$ .
3.  $\mathcal{A}$  gives access to its loss function  $\mathcal{L}$  and gradient of the loss function  $\nabla_{\rho}\mathcal{L}$  with respect to  $\rho$
These requirements are mildly restrictive and they are fulfilled by many gradient-based discovery methods (for instance, ENCO [Lippe et al., 2022], SDI [Ke et al., 2019], DiBS [Lorch et al., 2021], DCDI [Brouillard et al., 2020] or DECI [Geffner et al., 2022]).
Method. GIT scores each possible intervention target by calculating the expected magnitude of the gradient using imaginary interventional data generated by the causal model. Gradient magnitude serves as a proxy for the size of the update that can be induced on the parameters of the causal model.
The method picks intervention that has the highest score. Formally, for a given intervention  $i \in V$  we define its score  $s_i$  as follows:
$$
s _ {i} \triangleq \mathbb {E} _ {X \sim \mathbb {P} _ {\rho , \phi , i}} \| \nabla_ {\rho} \mathcal {L} (X) \|. \tag {4}
$$
Note that the expected value is computed with the interventional distribution coming from the model, instead of ground truth, defined as:
$$
\mathbb {P} _ {\rho , \phi , i} (X) \triangleq \sum_ {G} \mathbb {P} _ {\rho} (G) \mathbb {P} _ {G, \phi , i} (X). \tag {5}
$$
The summation in equation 5 is taken over all DAGs and  $\mathbb{P}_{G,\phi,i}$  corresponds to the joint distribution from the model for graph  $G$ :
$$
\mathbb {P} _ {G, \phi , i} (X) \triangleq \widetilde {\mathbb {P}} \left(X _ {i} \mid P A _ {i}\right) \prod_ {j \neq i} \mathbb {P} _ {G, \phi} \left(X _ {j} \mid P A _ {(j, G)}\right). \tag {6}
$$
The computational procedure of GIT's intervention target selection is listed in Algorithm 2. The expected value in  $s_i$  is approximated using the Monte-Carlo method, see line 4 of Algorithm 2. We also use a version of Algorithm 2 where real interventional data are used in line 3 (instead of the imaginary ones from the model) and call it GIT-privileged. GIT-privileged serves as a soft upper bound in our analysis. Note however, that it is not practically useful, as collecting real interventional data would require intervention on every node in the first place.
Algorithm 2 GIT'S INTERVENTION TARGET SELECTION
input parameters  $\rho$  of distribution over graphs, functional parameters  $\phi$ , loss function  $\mathcal{L}$ , graph nodes  $V$   
output batch of intervention targets to execute:  $I$
1:  $\mathcal{G}\gets$  sample a set of DAGs according to  $\mathbb{P}_{\rho}(G)$
2: for intervention target  $i \in V$  do
3:  $\mathcal{D}_{G,i}\gets$  sample batch of data according to  $\mathbb{P}_{G,\phi ,i}$
4:  $s_i \gets \frac{1}{|\mathcal{G}|} \sum_{G \in \mathcal{G}} \frac{1}{|\mathcal{D}_{G,i}|} \sum_{X \in \mathcal{D}_{G,i}} \| \nabla_\rho \mathcal{L}(X) \|$
6:  $I \gets$  select a batch of targets with highest scores  $s_i$
Assumptions of GIT. From the causal perspective, GIT relies exclusively on the Markov property assumption, which allows factorization of joined distribution (see Equation 1). However, GIT as a plug-and-play extension for causal discovery algorithms  $\mathcal{A}$  inherits their assumptions. This may include, for instance, causal sufficiency or faithfulness. Our method does not require any additional assumptions on the variables  $X_{i}$  and allows for both discrete and continuous setups.
Theoretical justification of GIT. We show the convergence of GIT in two contexts. First, we prove that the main setup of this paper, i.e., GIT with ENCO [Lippe et al., 2022], described in Section 4.1, converges. The detailed result can be found in Appendix B, but the gist of the argument is that vertices for which the model structure is not aligned with the ground truth will have non-trivial gradients and hence will be queried by GIT, allowing the model to improve. Moreover, we show empirically that GIT gradients are well correlated with the principled GPG signal of GIT-privileged, see Appendix F.6. Second, we show that given any convergent causal discovery algorithm, GIT converges if we allow a uniform sampling of intervention with small probability  $\epsilon > 0$ , see Appendix A. We call this approach  $\epsilon$ -greedy GIT. Importantly, on a finite sample with small enough  $\epsilon$ , GIT and  $\epsilon$ -greedy GIT are statistically indistinguishable.
4.1 Applicability to ENCO
We choose to use ENCO as the gradient-based causal discovery framework  $\mathcal{A}$  in our main experiments (recall Algorithm 1) due to its strong empirical results and good computational performance on GPUs. ENCO maintains a parameterized distribution over graph structures, with the so-called structural parameters  $\{\rho_{i,j}\}_{i,j}$  representing the adjacency matrix and a set of parameters modeling the functional dependencies,  $\phi$ . The structural parameters,  $\rho_{i,j}$ , are factorized into an edge existence parameter,  $\gamma_{i,j}$ , and an edge orientation parameter,  $\theta_{i,j} = -\theta_{j,i}$ .
The parameters are updated by iteratively alternating between two optimization phases. The goal of the first phase is to learn functions  $f_{\phi_i}(x_i|PA_{(i,G)})$ , which model the conditional density of  $\mathbb{P}\left(X_i|PA_{(i,G)}\right)$ . The training objective is the log-likelihood loss. The second phase aims to update the parametrized edge probabilities  $\rho_{i,j}$ 's. To this end, ENCO collects a data sample from a mixture
Table 1: We count the number of setups (24), where a given method was best or comparable to the other methods (AIT, CBED, Random, and GIT; GIT-privileged was not compared against), based on  $90\%$  confidence intervals for SHD and AUSHD. Each entry shows the total count, broken down into two data regimes,  $N = 1056$  and  $N = 3200$ , respectively, presented in parentheses.
AIT	CBED	Random	GIT (ours)	| GIT-privileged
mean AUSHD	6 (2 + 4)	6 (4 + 2)	12 (5 + 7)	18 (11 + 7)	24 (12 + 12)
mean SHD	10 (4 + 6)	7 (4 + 3)	22 (12 + 10)	17 (10 + 7)	24 (12 + 12)
of interventional distributions denoted by  $\mathbb{P}_I$ . The graph parameters are optimized by minimizing  $\mathbb{E}_{X\sim \mathbb{P}_I}L_{\mathrm{graph}}(X)$  where:
$$
L _ {\operatorname {g r a p h}} (X) \triangleq \mathbb {E} _ {G \sim P _ {\gamma , \theta}} \left[ \sum_ {i = 1} ^ {n} L _ {G} \left(X _ {i}\right) \right], \quad L _ {G} \left(x _ {i}\right) \triangleq - \log f _ {\phi_ {i}} \left(x _ {i} \mid P A _ {(i, G)}\right), \tag {7}
$$
For a detailed description of the method, distributions, and the estimators see Appendix C.1.
GIT with ENCO details. The loss function  $\mathcal{L}$  utilized by GIT is denoted  $L_{\mathrm{graph}}$ . We incorporate information from both structural parameters and use  $\| \nabla_{\gamma}L_{\mathrm{graph}}(X)\|^2 + \| \nabla_{\theta}L_{\mathrm{graph}}(X)\|^2$  to compute the score for the intervention  $i$  in line 4 of Algorithm 2. In order to sample DAGs from the current graph distribution (line 1 of Algorithm 2), we use a two-phase sampling procedure proposed in [Scherrer et al., 2021, Section 3.2] as it is scalable and guarantees DAG-ness by construction opposed to Gibbs sampling or rejection sampling techniques.
We compare GIT against the following baselines: AIT, CBED, Random, and GIT-privileged. AIT and CBED are competitive intervention acquisition methods for gradient-based causal discovery (which we discussed in Section 3.2). The Random method selects interventions uniformly in a round-robin fashion<sup>2</sup>. The last approach, GIT-privileged, is the oracle method described in Section 4.
Our main result is that GIT brings substantial improvement in the low data regime, being the best among benchmarked methods for all considered synthetic graph classes and half of the considered real graphs in terms of the AUSHD metric (see Equation 9). On the remaining real graphs, our approach performs similarly to the baseline methods. Notably, in most cases, GIT surpasses MI-based approaches: CBED and AIT. We present the summary in Table 1. This result is accompanied by an in-depth analysis of the relationships between different strategies and the distributions of the selected intervention targets. Additional results in the DiBS framework [Lorch et al., 2021] with continuous data are presented in Appendix F.1.