SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Identification: message passing formulation. Let  $\mathcal{I}$  be an IFM representing a model for the  $m$ -dimensional random vector  $X$  under the  $d$ -dimensional intervention vector  $\sigma \in \Sigma$ . Model  $\mathcal{I}$  has unknown factor parameters but a known factorization  $p(x; \sigma) \propto \prod_{k} f_{k}(x_{S_{k}}; \sigma_{F_{k}})$ .
Theorem 3.1 Assume that the  $\sigma$ -graph  $\mathcal{G}_{\sigma(\mathcal{I})}$  is decomposable. Given a set  $\Sigma_{\mathrm{train}} = \{\sigma^1, \ldots, \sigma^t\} \subseteq \Sigma$  and known distributions  $p(x; \sigma^1), \ldots, p(x; \sigma^t)$ , the following conditions are sufficient to identify any  $p(x; \sigma^\star)$ ,  $\sigma^\star \in \Sigma$ :
i) all distributions indexed by  $\sigma^1, \ldots, \sigma^t, \sigma^\star$  have the same support; and
ii)  $\Sigma_{[F_k]}^0\in \Sigma_{\mathrm{train}}$  for all  $F_{k}$  in the factorization of  $\mathcal{I}$
The algorithm for computing  $p(x; \sigma^{\star})$  works as follows. Construct a junction tree  $\mathcal{T}$  for  $\mathcal{I}$ , choose an arbitrary vertex in  $\mathcal{T}$  to be the root, and direct  $\mathcal{T}$  accordingly. If  $V_{k}$  is a hypervertex in  $\mathcal{T}$ , let  $D_{k}$  be the union of all intervention variables contained in the descendants of  $V_{k}$  in  $\mathcal{T}$ . Let  $B_{k}$  be the intersection of the intervention variables contained in  $V_{k}$  with the intervention variables contained in the parent  $V_{\pi(k)}$  of  $V_{k}$  in  $\mathcal{T}$ .
We define a message from a non-root vertex  $V_{k}$  to its parent  $V_{\pi (k)}$  as
$$
m _ {k} ^ {x} := \frac {p (x ; \sigma^ {[ D _ {k} (\star) ]})}{p (x ; \sigma^ {[ B _ {k} (\star) ]})}, \tag {2}
$$
with the update equation
$$
p (x; \sigma^ {[ D _ {k} (\star) ]}) \propto p (x; \sigma^ {[ F _ {k} (\star) ]}) \prod_ {V _ {k ^ {\prime}} \in c h (k)} m _ {k ^ {\prime}} ^ {x}, \tag {3}
$$
where the product over  $ch(k)$ , the children of  $V_{k}$  in  $\mathcal{T}$ , is defined to be 1 if  $ch(k) = \emptyset$ .
(All proofs in Appendix B.) In particular, if no factor contains more than one intervention variable, the corresponding  $\sigma$ -graph will be fully disconnected. This happens, for example, for an IFM derived from a DAG without hidden variables and where each intervention has one child, as in Fig. 2.
3.2 Algebraic Formulation
If a  $\sigma$ -graph is not decomposable, the usual trick of triangulation prior to clique extraction can be used [46, 20], at the cost of creating cliques which are larger than the original factors. Alternatively, and a generalization of Eqs. (2) and (3), we consider transformations of the distributions in  $\mathbb{P}(\Sigma_{\mathrm{train}})$
Quick summary: a junction tree is formed by turning the cliques of an undirected graph into the (hyper)vertices of the tree. Essential to the definition, a junction tree  $\mathcal{T}$  must have a running intersection property: given an intersection  $S\coloneqq H_{i}\cap H_{j}$  of any two hypervertices  $H_{i},H_{j}$ , all hypervertices in the unique path in  $\mathcal{T}$  between  $H_{i}$  and  $H_{j}$  must contain  $S$ . This property captures the notion that satisfying local agreements ("equality of properties" of subsets of elements between two adjacent hypervertices in the tree) should imply global agreements. A decomposable graph is simply a graph whose cliques can be arranged as a junction tree.
σ1	σ2	σ3	f100	f101	f110	f111	f200	f201	f210	f211	f300	f301	f310	f311
q1	0	0	0	✓	✓	✓
q2	0	1	0	✓	✓	✓
q3	1	0	0	✓	✓	✓
q4	1	1	0	✓	✓	✓
q5	0	0	1	✓	✓	✓
q6	0	1	1	✓	✓	✓
q7	1	0	1	✓	✓	✓
p*	1	1	1	0	0	0	1	0	0	0	1	0	0	0	1
Figure 4: Example of how to infer the target distribution  $p(x; (1, 1, 1))$  from the other seven possible regimes given the postulated factorization  $p(x; \sigma) \propto f_1(x; \sigma_1, \sigma_2)f_2(x; \sigma_2, \sigma_3)f_3(x; \sigma_1, \sigma_3)$ . We use  $f_k^{ab}$  as a shorthand notation for  $f_k(x; \sigma_{k_1} = a, \sigma_{k_2} = b)$ , e.g.,  $p(x; (1, 1, 1)) \propto f_1^{11}f_2^{11}f_3^{11}$ . For a PR-transformation  $(f_1^{00}f_2^{00}f_3^{00})^{q_1} \times (f_1^{01}f_2^{10}f_3^{00})^{q_2} \times \dots \times (f_1^{10}f_2^{01}f_3^{11})^{q_7}f_1^{11}f_2^{11}f_3^{11}$  up to a multiplicative constant independent of  $x$ , we need to find a solution satisfying:  $q_1 + q_5 = 0$ ,  $q_2 + q_6 = 0$ , ...,  $q_7 = 1$ . (In the table, this can be seen by going through each  $f_k^{ab}$  column and adding up the corresponding exponents  $q_j$  when there is a tick in row  $j$ ; the sum must agree with the corresponding entry in the final row). A solution is  $q_1 = q_4 = q_6 = q_7 = 1$ ,  $q_2 = q_3 = q_5 = -1$ , corresponding to  $p(x; \sigma^{\star}) \propto (p(x; \sigma^1)p(x; \sigma^4)p(x; \sigma^6)p(x; \sigma^7)) / (p(x; \sigma^2)p(x; \sigma^3)p(x; \sigma^5))$ .
by products and ratios. We define a PR-transformation of a density set  $\{p(x;\sigma^1),\ldots ,p(x;\sigma^t)\}$  as any formula  $\prod_{i = 1}^{t}p(x;\sigma^{t})^{q_{i}}$ , for a collection  $q_{1},\dots,q_{t}$  of real numbers.
Theorem 3.2 Let  $\sigma_{F_k}^v$  denote a particular value of  $\sigma_{F_k}$ , and let  $\mathbb{D}_k$  be the domain of  $\sigma_{F_k}$ . Given a collection  $\mathbb{P}(\Sigma_{\mathrm{train}}) \coloneqq \{p(x; \sigma^1), \ldots, p(x; \sigma^t)\}$  and a postulated model factorization  $p(x; \sigma) \propto \prod_{k=1}^{l} f_k(x_{S_k}; \sigma_{F_k})$ , a sufficient and almost-everywhere necessary condition for a given  $p(x; \sigma^\star)$  to be identifiable by PR-transformations of  $\mathbb{P}(\Sigma_{\mathrm{train}})$  is that there exists some solution to the system
$$
\forall k \in \{1, 2, \dots , l \}, \forall \sigma_ {F _ {k}} ^ {v} \in \mathbb {D} _ {k}, \left(\sum_ {i = 1: \sigma_ {F _ {k}} ^ {i} = \sigma_ {F _ {k}} ^ {v}} ^ {t} q _ {i}\right) = \mathbb {1} \left(\sigma_ {F _ {k}} ^ {\star} = \sigma_ {F _ {k}} ^ {v}\right), \tag {4}
$$
where  $\mathbb{1}(\cdot)$  is the indicator function returning 1 or 0 if its argument is true or false, respectively.
The solution to the system gives the PR-transformation. An example is shown in Fig. 4. The main idea is that, for a fixed  $x$ , any factor  $f_{k}(x_{S_{k}}; \sigma_{F_{k}})$  can algebraically be interpreted as an arbitrary symbol indexed by  $\sigma_{F_{k}}$ , say " $f_{k}^{\sigma_{F_{k}}}$ ". This means that any density  $p(x; \sigma^{i})$  is (proportional to) a "squarefree" monomial  $m^{i}$  in those symbols (i.e., products with exponents 0 or 1). We need to manipulate a set of monomials (the unnormalized densities in  $\mathbb{P}(\Sigma_{\mathrm{train}})$ ) into another monomial  $m^{\star}$  (the unnormalized  $p(x; \sigma^{\star})$ ). More generally, the possible set functions  $g(\cdot), h(\cdot)$  such that  $g(m^{\star}, \{m^{1}, \dots, m^{t}\}) = h(\{m^{1}, \dots, m^{t}\})$ , with  $g(\cdot)$  invertible in  $m^{\star}$ , suggest that  $g(\cdot)$  and  $h(\cdot)$  must be the product function (monomials are closed under products, but not other analytical manipulations), although we stop short of stating formally the conditions required for PR-transformations to be complete in the space of all possible functions of  $\mathbb{P}(\Sigma_{\mathrm{train}})$ . This would be a stronger claim than Theorem 3.2, which shows that Eq. (4) is almost-everywhere complete in the space of PR transformations.
The message passing scheme and the algebraic method provide complementary views, with the former giving a divide-and-conquer perspective that identifies subsystems that can be estimated without requiring changes from the baseline treatment everywhere else. The algebraic method is more general, but suggests no hierarchy of simpler problems. These results show that the factorization over  $X$  is unimportant for identifiability, which may be surprising. Appendix C discusses the consequences of this finding, along with a discussion about requirements on the size of  $\Sigma_{\mathrm{train}}$ .
4 Learning Algorithms
The identification results give a license to choose any estimation method we want if identification is established — while we must make the assumption of the model factorizing according to a postulated IFM, we are not required to explicitly use a likelihood function. Theorems 3.1 and 3.2 provide ways of constructing a target distribution  $p(x; \sigma^{\star})$  from products of ratios of densities from  $\mathbb{P}(\Sigma_{\mathrm{train}})$ . This
suggests a plug-in approach for estimation: estimate products of ratios using density ratio estimation methods and multiply results. However, in practice, we found that fitting a likelihood function directly often works better than estimating the product of density ratios, even for intractable likelihoods. We now discuss three strategies for estimating some  $\mathbb{E}[Y; \sigma^{\star}]$  of interest.
Deep energy-based models and direct regression. The most direct learning algorithm is to first maximize the sum of log-likelihoods  $\mathcal{L}(\theta; \mathcal{D}^1, \ldots, \mathcal{D}^t) := \sum_{i=1}^{t} \sum_{j=1}^{n_i} \log p_\theta(x^{i(j)}; \sigma^i)$ , where  $n_i$  is the number of samples in the  $i$ th regime. The parameterization of the model is indexed by a vector  $\theta$ , which defines  $p(\cdot)$  as  $\log p_\theta(x; \sigma) := \sum_{k=1}^{l} \phi_{\theta_k, \sigma_{F_k}}(x_{S_k}) + \text{constant}$ . Here,  $\phi_\theta(\cdot)$  is a differentiable black-box function, which in our experiments is a MLP (aka a multilayer perceptron, or feedforward neural network). Parameter vector  $\theta_k, \sigma_{F_k}$  is the collection of weights and biases of the MLP, a different instance for each factor  $k$  and combination of values in  $\sigma_{F_k}$ . In principle, making the parameters smooth functions of  $\sigma_{F_k}$  is possible, but in the interest of simplifying the presentation, we instead use a look-up table for completely independent parameters as indexed by the possible values of  $\sigma_{F_k}$ . Parameter set  $\theta$  is the union of all  $\sum_{k=1}^{l} \prod_{j \in F_k} |\aleph_j|$  MLP parameter sets. As maximizing log-likelihood is generally intractable, in our implementation we apply pseudo-likelihood with discretization of each variable  $X$  in a grid. The level of discretization does not affect the number of parameters, as we take their numerical value as is, renormalizing over the pre-defined grid. Score matching [38], noise contrastive estimation [32] or other variants (e.g., [71]) could be used; our pipeline is agnostic to this choice, with further details in Appendix G. We then estimate  $f_y(x) := \mathbb{E}[Y | x]$  using an off-the-shelf method, which in our case is another MLP. For a given  $\sigma^\star$ , we sample from the corresponding  $p(x; \sigma^\star)$  with Gibbs sampling and average the results of the regression estimate  $\widehat{f}_y(x)$  to obtain an estimate  $\widehat{\mu}(\sigma^\star)$ .
Inverse probability weighting (IPW). A more direct method is to reweight each training sample by the target distribution  $p(x; \sigma^{\star})$  to generate  $\widehat{\mu}(\sigma^{i\star}) := \sum_{j=1}^{n_i} y^{i(j)} w^{i(j)^{\star}}$ , where  $w^{i(j)^{\star}}$  is the density ratio  $p(x^{i(j)}; \sigma^{\star}) / p(x^{i(j)}; \sigma^i)$ , rescaled such that  $\sum_{j=1}^{n_i} w^{i(j)^{\star}} = 1$ . There are several direct methods for density ratio estimation [53] that could be combined using the messages/product-ratios of the previous section, but we found that it was stabler to just take the density ratio of the fitted models using deep energy-based learning, just like in the previous algorithm. Once estimators  $\widehat{\mu}(\sigma^{\star(1)}), \ldots, \widehat{\mu}(\sigma^{\star(t)})$  are obtained, we combine them by the usual inverse variance weighting rule,  $\widehat{\mu}(\sigma^{i\star}) := \sum_{i=1}^{t} r^i \mu_{\sigma^{\star(i)}} / \sum_{i=1}^{t} r^i$ , where  $r^i := 1 / \widehat{v}^i$ , and  $\widehat{v}^i := \sum_{j=1}^{n_i} (y^{i(j)})^2 (w^{i(j)^{\star}})^2$ . This method requires neither a model for  $f_y(x)$  nor Markov chain Monte Carlo. However, it may behave more erratically than the direct method described above, particularly under strong shifts in distribution.
Covariate shift regression. Finally, a third approach for estimating  $\mu (\sigma^{\star})$  is to combine models for  $f_{y}(x)$  with density ratios, learning a customized  $\widehat{f}_y(x)$  for each test regime separately. A s this is very slow and did not appear to be advantageous compared to the direct method, we defer a more complete description to Appendix D.
Predictive Coverage. Even when treatment effects are identifiable within the IFM framework, the uncertainty of resulting estimates can vary widely depending on the training data and learning algorithm. Building on recent work in conformal inference [78, 51, 76], we derive the following finite sample coverage guarantee for potential outcomes, which requires no extra assumptions beyond those stated above.
Theorem 4.1 (Predictive Coverage.) Assume the identifiability conditions of Theorems 3.1 or 3.2 hold. Fix a target level  $\alpha \in (0,1)$ , and let  $\mathcal{I}_1, \mathcal{I}_2$  be a random partition of observed regimes into training and test sets of size  $n/2$ . Fit a model  $\widehat{\mu}$  using data from  $\mathcal{I}_1$  and compute conformity scores  $s^{(i)} = |y^{k(i)} - \widehat{\mu}(\sigma^k)|$  using data from  $\mathcal{I}_2$ . For some new test environment  $\sigma^\star$ , compute the normalized likelihood ratio  $w^{(i)}(\sigma^\star) \propto p(x^{k(i)}; \sigma^\star)/p(x^{k(i)}; \sigma^k)$ , rescaled to sum to  $n/2$ . Let  $\widehat{\tau}(\sigma^\star)$  be the  $q^{\text{th}}$  smallest value in the reweighted empirical distribution  $\sum_{i} w^{(i)}(\sigma^\star) \cdot \delta(s^{(i)})$ , where  $\delta$  denotes the Dirac delta function and  $q = \lceil (n/2 + 1)(1 - \alpha) \rceil$ . Then for any new sample  $n + 1$ , we have:
$$
\mathbb {P} \left(Y ^ {\star (n + 1)} \in \widehat {\mu} \left(\sigma^ {\star}\right) \pm \widehat {\tau} \left(\sigma^ {*}\right)\right) \geq 1 - \alpha .
$$
Moreover, if weighted conformity scores have a continuous joint distribution, then the upper bound on this probability is  $1 - \alpha + 1 / (n / 2 + 1)$ .
We run a number of semi-synthetic experiments to evaluate the performance of the IFM approach on a range of intervention generalization tasks. In this section, we summarize our experimental set-up and main results. The code for reproducing all results and figures is available online<sup>8</sup>; in Appendix E, we provide a detailed description of the datasets and models; and in Appendix F we present further analysis and results.
Datasets. Our experiments are based on the following two biomolecular datasets:  $i$  Sachs [66]: a cellular signaling network with 11 nodes representing phosphorylated proteins and phospholipids, several of which were perturbed with targeted reagents to stimulate or inhibit expression. There are 4 binary intervention variables.  $\Sigma_{\mathrm{train}}$  has 5 regimes: a baseline, with all  $\sigma_{i} = 0$ , and 4 "single-intervention" regimes, each with a different single  $\sigma_{i} = 1$ .  $\Sigma_{\mathrm{test}}$  consists of the remaining 11  $(= 2^{4} - 5)$  unseen regimes. ii) DREAM [30]: Simulated data based on a known E. coli regulatory sub-network with 10 nodes. There are 10 binary intervention variables, and (similarly)  $\Sigma_{\mathrm{train}}$  has a baseline regime, with all  $\sigma_{i} = 0$ , and 10 single-intervention regimes, with single choices of  $\sigma_{i} = 1$ .  $\Sigma_{\mathrm{test}}$  consists of 45  $(10 \times 9/2 = 45)$  unseen regimes defined by all pairs  $\sigma_{i} = \sigma_{j} = 1$ .
Oracular simulators. The first step to test intervention generalization is to build a set of proper test beds (i.e., simulators that serve as causal effect oracles for any  $\sigma \in \Sigma$ ), motivated by expert knowledge about the underlying system dynamics. Neither the original Sachs data nor the DREAM simulator we used provide joint intervention data required in our evaluation. Thus, for each domain, we trained two ground truth simulators:  $i)$  Causal-DAG: a DAG model following the DAG structure and data provided by the original Sachs et al. and DREAM sources (DAGs shown in Fig. 10(a)) and Fig. 11(a), respectively). Given the DAG, we fit a model where each conditional distribution is a heteroskedastic Gaussian with mean and variance parameterized by MLPs (with 10 hidden units) of the respective parents.  $ii)$  Causal-IFMs: the corresponding IFM is obtained by a direct projection of the postulated DAG factors (as done in, e.g., Fig. 2(b)). The likelihood is a neural energy model (Section 4) with MLPs with 15 hidden units defining potential functions. After fitting these models, we compute by Monte Carlo simulation their implied ground truths for every choice of  $\sigma^{*}$ . The outcomes  $Y$  are then generated under 100 different structural equations of the form  $\tanh (\lambda^{\top}X) + \epsilon$ , with random independent normal weights  $\lambda$  and  $\epsilon \sim \mathcal{N}(0,v_y)$ .  $\lambda$  and  $v_{y}$  are scaled such that the ground truth variance of  $\lambda^{\top}X$ ,  $\operatorname{Var}(\lambda^{\top}X)$ , is sampled uniformly at random from the interval [0.6, 0.8], and set  $v_{y} := 1 - \operatorname{Var}(\lambda^{\top}X)$ .
Compared models. We implement three variants of our proposed IFM model, corresponding to the three learning algorithms described in Section 4: i) IFM1 uses deep energy-based models and direct regression; ii) IFM2 uses an IPW estimator; and iii) IFM3 relies on covariate shift regression. We compare these models to the following benchmarks: i) Black-box: We apply an off-the-shelf algorithm (XGBoost [14]) to learn a direct mapping from  $\sigma$  to  $Y$  without using  $X$ . This model does not exploit any structural assumptions. ii) DAG: We estimate the structural equations in an acyclic topological ordering that is consistent with the data generating process. This represents a strong baseline that exploits ground truth knowledge of the underlying causal graph. The likelihood is defined by conditional Gaussian models with mean and variances parameterized as MLPs, matching exactly one of the simulators described below.
Results. We evaluate model performance based on the proportional root mean squared error (pRMSE), defined as the average of the squared difference between the ground truth  $Y$  and estimated  $\bar{Y}$ , with each entry further divided by the ground truth variance of the corresponding  $Y$ . Results are visualized in Fig. 5. We additionally run a series of one-sided binomial tests to determine whether models significantly outperform the black box baseline, and compare the Spearman's rank correlation [80] between expected and observed outcomes for all test regimes. Unsurprisingly, DAGs do best when the ground truth is a Causal-DAG, while IFM methods do better when the data generator is a Causal-IFM. Still, some IFMs are robust to both ground truth models, with IFM1 (the deep energy model) doing especially well on the DREAM dataset, and IFM2 (the IPW estimator) excelling on the Sachs data.
<sup>8</sup>https://github.com/rbas-ucl/intgen
Figure 5: Experimental results on a range of intervention generalization tasks, see text for details.
In particular, for the DREAM datasets, IFM1 (scaled) errors remain stable, while the DAG has a major loss of performance when a non-DAG ground truth is presented. However, IFM2 underperform on DREAM datasets, which might be because the knockdowns regimes in those datasets induce dramatic shifts in distribution overlap among regimes. (We omit results for IFM3 (covariate shift regression) on these datasets, as the method does not converge in a reasonable time.) Though the black-box method (XGBoost algorithm and no structural assumption) sometimes struggles to extrapolate, displaying a long-tailed distribution of errors, it actually does surprisingly well in some experiments, especially on the DREAM dataset. In fact, the black-box method is essentially indistinguishable from linear regression in this benchmark (results not shown), which is not surprising given the sparsity of the training data. However, the non-additivity of the  $\sigma$  effect on  $Y$  is more prominent in the Sachs datasets, making it difficult to generalize without structural assumptions.
A factor graph interpretation of Pearl's do-calculus framework is described in [79], but without addressing the problem of identifiability. More generally, several authors have exploit structural assumptions to predict the effects of unseen treatments. Much of this research falls under the framework of transportability [61, 7], where the goal is to identify causal estimands from a combination of observational and experimental data collected under different regimes. If only atomic interventions are considered, the do-calculus is sound and complete for this task [50]. The  $\sigma$ -calculus introduced by [28, 17] extends these transportability results to soft interventions [18, 19]. However, these methods are less clearly defined when an intervention affects several variables simultaneously, and while DAGs with additive errors have some identifiability [67] and estimation [75] advantages, acyclic error additivity may not be an appropriate assumption in some domains. Under further parametric assumptions, causal effects can be imputed with matrix completion techniques [73, 3] or more generic supervised learning approaches [31], but these methods often require some data to be collected for all regimes in  $\Sigma_{\mathrm{test}}$ . Finally, [2] is the closest related work in terms of goals, factorizing the function space of each  $\mathbb{E}[X_i; \sigma]$  directly as a function of  $\sigma$ .
Another strand of related research pertains to online learning settings. For instance, several works have shown that causal information can boost convergence rates in multi-armed bandit problems when dependencies are present between arms [45, 49, 24], even when these structures must themselves be adaptively learned [9]. This suggests a promising direction for future work, where identification strategies based on the IFM framework are used to prioritize the search for optimal treatments. Indeed, combining samples across multiple regimes can be an effective strategy for causal discovery, as illustrated by recent advances in invariant causal prediction [63, 33, 64, 81], and it can also help with domain adaptation and covariate shift [54, 6, 13, 15].
We introduced the IFM framework for solving intervention generalization tasks. Results from simulations calibrated by real-world data show that our method successfully predicts outcomes for novel treatments, providing practitioners with new methods for conducting synthetic experiments. Future work includes:  $i$  ) integrating the approach with experimental design, Bayesian optimization and bandit learning; ii) variations that include pre-treatment variables and generalizations across heterogeneous subpopulations, inspired by complementary matrix factorization methods such as [2]; iii) tackling sequential treatments; and iv) diagnostics of cross-regime overlap issues [36, 58].
We thank the anonymous reviewers and Mathias Drton for useful discussions. GBH was supported by the ONR grant 62909-19-1-2096, and JY was supported by the EPSRC grant EP/W024330/1. RS was partially supported by both grants. JZ was supported by UKRI grant EP/S021566/1.