SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For the case where  $\Pi^*$  is not known, let  $\widehat{\Pi}$  be the set of recovered partitions for each  $X_{i}$ . Then if  $\widehat{\Pi} \neq \Pi^*$  there must be at least one  $i$  for which  $\widehat{\Pi}_i \neq \Pi_i^*$ . There are two ways in which the partitions can differ.
First, the partition  $\widehat{\Pi}_i$  is more fine-grained than  $\Pi_i^*$ . That is, for each  $\pi^* \in \Pi_i^*$  there exist groups  $\pi_1, \ldots, \pi_k \in \widehat{\Pi}_i$  such that  $\pi^* = \pi_1 \cup \ldots \cup \pi_k$  and further there exists at least one  $\pi^*$  for which  $k > 1$ . For this  $\pi^*$  and associated  $\pi_1, \ldots, \pi_k$  we have as before that for sufficiently small  $\lambda$  the minimizer for  $L(X; \widehat{G})$  would have to be  $f_i^{\pi_1} = \ldots = f_i^{\pi_k}$ . However, by assumption (A3) this is not permitted. This case therefore cannot occur.
Second, there could be a group  $\pi \in \widehat{\Pi}_i$  which intersects at least two different groups  $\pi_1^* \neq \pi_2^* \in \Pi_i^*$ . However, from (A3) we know that  $f_i^{\pi_1^*} \neq f_i^{\pi_2^*}$  so that whatever function  $f_i^\pi$  we pick in the discovered model, it would necessarily lead to the wrong distribution for at least one of the two groups and therefore a positive  $D_{\mathrm{KL}}$  term, which as before is sub-optimal for  $\lambda$  small enough.
For the proof of Theorem 3.2, we require the following formalization of independent changes of mechanisms.
Assumption 5 (Independent Mechanism Shift): For all  $i \neq j: \widehat{\Pi}_i(C) \perp \widehat{\Pi}_j(C)$ .
Under this assumption, we can prove that we can identify the full causal model.
Theorem 3.2 (Identifiability from independent changes): Let Assumptions 1-5 hold. Then for  $\lambda$  sufficiently small, where  $\lambda$  is a hyperparameter bounding the score differences of  $\widehat{G}, G^{*}$ , we can identify the causes  $pa_{i}^{G}$  and partition  $\Pi_{i}$  of each  $X_{i}$  in the causal model  $(G^{*}, \Pi^{*})$  with probability 1,
Proof. As we have seen from Theorem 3.1, we can identify the correct MEC and  $\Pi^*$  already. Our goal is therefore to show that any other  $\widehat{G}$  in the same MEC as  $G^*$  will correspond to a  $\widehat{\Pi}$  different from  $\Pi^*$ .
Let us consider the target variable  $X_{i}$ . We can ask ourselves what happens when an edge  $X_{i}\to X_{j}$  is misdirected in  $\widehat{G}$ . We will show that the misdirected edge will lead to a wrong partition with probability 1 as  $|C|\rightarrow \infty$
To that end, note that when  $\Pi_i^*(c) = \Pi_i^*(c')$  but  $\Pi_j^*(c) \neq \Pi_j^*(c')$  for two different contexts  $c, c'$  then in the alternative graph  $\widehat{G}$ , the distribution  $P(X_i \mid \mathrm{pa}_i^{\widehat{G}}, C = c) \neq P(X_i \mid \mathrm{pa}_i^{\widehat{G}}, C = c')$  is not invariant. Therefore, for each set  $\pi \in \widehat{\Pi}_i$ , the number of functions required to model the distribution of  $X_i$  increases. That is, we require one function for each  $\widehat{\Pi}_j(c)$  such that  $\widehat{\Pi}_i(c) = \pi$ .
Now, due to  $\widehat{\Pi}_i(C) \perp \widehat{\Pi}_j(C)$  this is guaranteed to be one function for each pair  $(\pi, \pi') \in \widehat{\Pi}_i \times \widehat{\Pi}_j$  as  $|\mathcal{C}| \to \infty$ .
However, this would lead to the wrong partitions  $\widehat{\Pi}_j \neq \Pi_j^*$ , in contradiction to Theorem 3.1 above.
A.4 Proof of Theorem 3.3
We now turn to the probabilistic version of the previous result, for which we need the following assumption.
Assumption 6 (Fixed Partition Sizes): For each  $X_{i}$ , there exists a finite number of sets  $B_{i,1},\ldots ,B_{i,k_i}$  such that  $p_{ij}\coloneqq P(C\in B_{i,j}) > 0$ .
With this, we can prove Theorem 3.3.
Theorem 3.3 (Identifiability in the finite context setting): Let the Assumptions of Theorem 3.2 as well as Assumption 6 hold. Let  $p = \min_{i,j} p_{ij}$  be the minimum probability for a context to fall into one of the bins  $B_{i,j}$ . Then if  $|\mathcal{C}| \geq -2p^{-2}\log(p) > 0$  we have
Proof. The proof is based on the same idea as the proof of Theorem 3.2, and uses probabilistic bounds on the coupon collector problem to obtain the desired result.
That is, we derive bounds on the probability that for an edge  $X_{i} \to X_{j}$  in  $G^{*}$ , there exist pairs of contexts  $c, c'$  such that  $\Pi_{i}^{*}(c) = \Pi_{i}^{*}(c')$  and  $\Pi_{j}^{*}(c) \neq \Pi_{j}^{*}(c')$  for all possible combinations in  $\widehat{\Pi}_i \times \widehat{\Pi}_j$ . We can consider these as bins with probabilities  $\geq p^2$ , so that for the stopping time  $T = |C|_{\mathrm{fin}}$  at which each such bin has been hit at least once we have
$$
P \left(T <   - 2 p ^ {- 2} \log (p) + c / p\right) \longrightarrow \exp (- e ^ {- c}).
$$
By using  $c = p^2 |\mathcal{C}| + 2\log p \geq 0$ , we therefore have
$$
P (X _ {i} \to X _ {j} \in \widehat {G}, \widehat {\Pi} _ {i} = \Pi_ {i} ^ {*}) = \Omega \left(\exp \left(- e ^ {- c}\right)\right),
$$
and by taking the naive lower bound of treating all edges independently, we therefore obtain
$$
P (\widehat {G} = G ^ {*}, \widehat {\Pi} = \Pi^ {*}) = \Omega \left(\exp \left(- | E | e ^ {- c}\right)\right).
$$
A.5 Discussion of Identifiability Results
Here, we place our identifiability results in the context of related work. In doing so, we explain our observation that our approach already discovers causal directions given a single context, as well as state the implications of Thm. 3.3 on how many contexts we need to identify the causal DAG.
Identifiability from a single context The work of Peters et al. [45] shows that causal DAGs are fully identifiable from i.i.d. data using independence of residuals. These results are obtained by assuming differentiability of the underlying distributions as well as causal mechanisms. By modeling our causal mechanisms via GPs, these assumptions are not applicable since individual GP paths are in general not differentiable. For our more general setting, we are not aware of a proof that the DAG is identifiable from a single context. However, Marx and Vreeken [25] show that in the bivariate case, MDL-based score functions allow to identify the causal direction from i.i.d. data using the independent noise asymmetry. Our empirical results support this as our approach can already direct a reasonable number of edges given a single context as well as in the bivariate case.
Identifiability depending on the number of contexts We remark that our results address the case where our data comes from random environments over which we have no control. When we do have control over the environments, i.e., when we can actively perform interventions with known effects, then general results on the numbers of required contexts apply, requiring data from  $\mathcal{O}(\log m)$  environments for  $m$  variables [32].
As a second note, in Theorem 3.3, we show how many contexts it takes us to identify both the causal graph as well as the partition of the causal mechanisms across environments. If we care only about the identifiability of the causal graph, then two mechanisms per variable (for those variables that are not already directed by Meek's rules) suffice, which dramatically reduces the number of contexts required. In particular, if we obtain samples with equal probabilities for each mechanism, then we require only 20 contexts to orient up to 10000 edges with probability  $>0.99$ , independently of the precise structure of the graph. This compares to (roughly) 8 contexts that are required under optimal intervention design for causal learning [32]. Our results however do not require us to have control over the intervention targets, which is realistic in many practical applications where we may encounter unknown distribution shifts or off-target intervention effects [7, 17].
A.6 Discussion of Assumptions
To conclude our theoretical analysis and point out the limitations of our method, we revisit the main assumptions that we make and discuss in which settings we consider them reasonable.
Functional Model (Assumption 1,2) While we avoid some issues related to conditional independence testing in causal discovery [14, 15], this comes at the cost that we need to make additional assumptions on the functional form of causal mechanisms. As we noted in the main text, Gaussian processes constitute a rich function class and allow for a fully non-parametric approach. While they are more general than the functional model assumed in Peters et al. [45], a similar kernelized scoring function is proposed by Huang et al. [15] which also allows for regularization, so a use of this functional model in a non-i.i.d. setting could be an interesting comparison in future work.
Invariant Noise (Assumption 1,2) By assuming that the noise  $N_{i}$  remains invariant across environments, we follow how related work defines environments [13, 9, 7]. This assumes that different environments arise from changes in the functional dependencies between the variables, commonly known as soft interventions. That is, environments differ due to changes of the conditional  $P(X_{i} \mid \mathrm{pa}_{i}^{G})$  of some variables  $X_{i}$  while the exogenous noise variables  $N_{i}$  remain invariant. Example applications where this holds include knockout or knockdown interventions in genomics, where the expression of a certain gene is prevented or inhibited [16, 17], as well as the molecular interventions in our real-world dataset, where activating or inhibiting reagents are added to the system [20].
To give an example where noise invariance is violated, we consider direct interventions on the noise  $N_{i}$  in our experiments (Fig. 3(a)). Our results suggest that this particular model misspecification is not a limitation in practice. The same is true for perfect intervention (Fig. 3(a)).
Sufficiency There is one other reason why noise invariance may be violated which relates to another assumption, namely causal sufficiency. If there is an unobserved variable, it may introduce a correlation of the noise  $N_{i}$  of different observed variables. Considering contexts allows to address confounding in a limited sense, since the categorical context variable  $C$  can represent so-called "pseudo-confounders", such as background stimuli that remain fixed in a given context [4]. However, these need to be a function of the domain index. There may also be a global latent confounder that introduces a correlation between the noises  $N^{(c)}$  across different environments. We are not aware of approaches that address this situation, and extending our theoretical framework to general latent confounding is an interesting option for future work.
Faithfulness (Assumption 3.4) Faithfulness assumptions are standard in causal discovery, and in our case assert that the causal parents of a variable indeed affect the variable, as well as that our partitions faithfully represent all changes in causal mechanisms. Both criteria are necessary to detect parents, respectively mechanism changes, from the observed distributions.
Independent Mechanism Shift (Assumption 5) Rather than assuming sparse mechanism shifts, we consider independent mechanism shifts in our work. To formally assume this, we consider our partitions  $\Pi$  as independent random variables, which can be thought of as deterministic functions of the context variable  $C$ . The result is a categorical random variable that indicates the group membership, rather than context membership, of samples. By assuming independence of these random variables and including them in our model, we ensure that incorrectly directed edges in  $G$  lead to larger numbers of modeled mechanisms and thus larger values of our score  $L$ .
To explain this, we recall the example in Fig. 1. The true partition for  $Y$  has two groups, and the partition of its causal child  $Z$  has either two (Fig. 1(b)) or five (Fig. 1(c)) groups. In the true causal model, the partitions for  $Y$  and  $Z$  are independent. Indeed, in the causal direction, we see no differences in our MDL scores between both cases (green curves in Fig. 1(b), 1(c)) and discover the true partition for  $Y$  with  $|f| = 2$  (green star) independently of the changing partition for  $Z$  in these cases. If we consider an anticausal function  $Z \rightarrow Y$ , in contrast, there is a difference between both cases (red curves). The partition we discover for  $Y$  here depends on the partition of  $Z$ , and in both cases it has more groups (minima of the red curves,  $|f| = 4$ ,  $|f| = 5$ ) and larger score values compared to the causal direction. In the proof of Theorem 3.2, we show more generally that under independence of the true partitions, incorrectly directed edges will lead to more groups in the discovered partitions as well as larger score values of the model.
We justify the independent mechanism shift assumption in the context of interventional experiments. Consider a gene knockout intervention. This is an exogenous influence and independent of the causal interactions between other genes. Similarly, Sachs et al. [20] perform soft interventions by adding compounds that activate or inhibit specific molecules. In both examples, a mechanism change is due to an external influence and does not depend on the causal mechanism of other variables. In general, this independence holds in generic situations where there is no confounding behind cause and effect [1, 4].
Fixed Partition Sizes (Assumption 6) By fixing partition sizes, we fix the number of different causal mechanisms for each of the variables. In effect, we assume that all heterogeneity in causal effects is captured in our model. For example, if we studied the differences between treatments in different hospitals, then as we obtain more data, these data would be obtained from the same set of hospitals, rather than by the addition of new hospitals. By ensuring that the partition sizes are fixed, we ensure that the things we are trying to learn do not change as we obtain more data. This allows us to give our more practical identifiability result in Theorem 3.3 for a smaller number of contexts.
In this section, we describe the implementation details of the LINC algorithm and its variants.
B.1 LINC for Learning causal models under Independent Changes
Our approach LINC considers data over continuous variables  $X^{(c)}$  in multiple contexts  $c \in \mathcal{C}$ , and returns a causal graph  $G$  and a set of partitions  $\Pi$ .
Algorithm 1: LINC
Input: data  $X^{(c)}$  , candidate DAGs G.   
Output: DAG  $G$  and partitions  $\Pi$    
for each DAG  $G\in \mathbf{G}$  do for each variable  $X_{i}$  with causes  $X_{S}$  in  $G$  do for each partition  $\Pi$  of  $C$  do for each group  $\pi$  in  $\Pi$  do Pool the data in  $\pi$  to obtain  $X^{(\pi)}$  Fit a GP  $f\sim \mathcal{GP}(X_S^{(\pi)}\to X_i^{(\pi)})$  Compute  $L(X_{S}^{(\pi)},X_{i}^{(\pi)}|\mathcal{H}_{\kappa})$  using Eq. (6) end for end for end for Find the partitions  $\widehat{\Pi_G}$  minimizing Eq. (10)   
end for   
Find the DAG  $\widehat{G}$  minimizing Eq. (11)   
return  $\widehat{G},\widehat{\Pi_G}$
In Alg. 1, we show the exact search algorithm that we consider in the main paper. It traverses a set of candidate causal graphs  $\mathbf{G}$ . For each graph  $G$  and variable  $X_{i}$ , it discovers a mixture of mechanisms that generate  $X_{i}$  from its parents in  $G$ . To do this, we consider each partition of the contexts, fit a GP regression model to the data pooled in each group in the partition, and compute the MDL score for these GPs. We choose the partition that results in the smallest overall description length,
$$
\begin{array}{l} \widehat {\Pi_ {i}} = \arg \min _ {\Pi} L (\Pi) \\ = \arg \min  _ {\Pi} \sum_ {\pi \in \Pi} L \left(X _ {S} ^ {(\pi)}, X _ {i} ^ {(\pi)} \mid \mathcal {H} _ {\kappa}\right) \tag {10} \\ \end{array}
$$
where the set  $X_{S}$  corresponds to the parents of  $X_{i}$  in  $G$ . Repeating this for each variable, we obtain the partitions  $\widehat{\Pi_G}$  for  $G$ . We return the DAG  $\widehat{G}$  with smallest overall description length,
$$
\widehat {G} = \arg \min  _ {G \in \mathbf {G}} \sum_ {\Pi \in \widehat {\Pi_ {G}}} L (\Pi). \tag {11}
$$
Regarding the complexity of our algorithm, we need to address three aspects: the search space over DAGs  $\mathbf{G}$ , the search space over partitions  $\Pi$ , and the GP regression itself. We propose efficient approaches for each aspect in the following, and conclude by giving a complexity analysis.
B.2 LINC with Greedy Equivalence Search
We first address the problem of discovering causal DAGs. As the search space over DAG models is super-exponential in the number of variables [33], we limit consideration to the Markov Equivalence class of the true model and set  $\mathbf{G} = \mathrm{MEC}(G^{*})$ . We remark that in case the true MEC is unknown, one can apply LINC using a two-stage approach, first using existing causal discovery approaches for multiple contexts [4, 7, 5, 9] to discover a partially directed causal graph, and subsequently applying LINC to orient the remaining edges.
We can also use our MDL score in a score-based DAG search algorithm, such as Greedy Equivalence Search (GES) [33]. GES uses a scoring criterion that evaluates each causal relationship between a given set of causes  $X_{S}$  and an effect  $X_{i}$ . Huang et al. [15] show how to use generalized, nonparametric scoring functions in GES, leading to the Generalized GES algorithm (GGES). Our score can be used as a plug-in estimator within this framework because it is a locally consistent scoring criterion due to information consistency of GPs. We note that our score is not score-equivalent, that is, different DAGs in the same MEC may attain different scores, but Huang et al. [15] show that score local consistency is sufficient to use a scoring function with GGES. While we discover a single DAG  $G$  over all contexts in this manner, we can account for non-i.i.d.-ness of the data as well as changes of the causal mechanisms.
Algorithm 2: LINCCLUS
Input: data  $X^{(c)}$  , variable  $X_{i}$  with causes  $X_{S}$    
Output: partition  $\Pi_{i}$    
for each context  $C$  in  $C$  do Regression in the context  $f^{(c)}\sim \mathcal{GP}(X_S^{(c)}\to X_i^{(c)})$  for each context  $c^{\prime}\neq c$  and  $c^\prime$  already seen do Compute  $L^{\prime}(X_S,X_i\mid c,c^{\prime})$  using Eq. (12) end for   
end for   
Find  $\Pi = \arg \min_{\Pi}\sum_{\pi \in \Pi}\sum_{c,c'\in \pi}L'(X_S,X_i\mid c,c')$  using clustering, e.g. by solving Eq. (13)   
return  $\Pi$
We propose exploring combinations of LINC with different DAG search frameworks in future work, and focus on orienting edges within a MEC in our experiments. This is motivated by two reasons. Firstly, it allows to study the behavior of our scoring function in isolation, taking out the effect of greedy DAG search algorithms, such that our experimental evaluation can support our theoretical results. Secondly, our closest competitors are the mechanism counting scores that need to start from a base MEC [12, 13] such that we also evaluate LINC on orienting edges within a MEC.
B.3 LINC with Clustering
To scale efficiently with the number of contexts, we introduce a variant of LINC that, rather than all set partitions, considers pairwise combinations of contexts.
To this end, we adapt the Minimum Description Length (MDL) score for GPs. We want to measure whether the data  $X^{(c)}$ ,  $X^{(c')}$  over two contexts  $c, c'$  is preferably encoded using separate GP functions for each, or a single model over the pooled data. A natural way to compare different modelling hypotheses is via the gain in compression of one model over the other.
By the no-hypercompression inequality [29], the probability that a model other than the optimal one should achieve a substantially better compression for a dataset is negligible. More specifically, the probability that the competing model leads to a gain in compression of  $k$  bits if it is not the true model is smaller than or equal to  $2^{-k}$ . For a significance threshold  $\alpha$ , we can consider a gain of  $k$  bits as insignificant in case  $2^{-k} > \alpha$  [29].
Hence, we consider the gain  $L'$  in compression that two separate GPs achieve over a joint model,
$$
\begin{array}{l} L ^ {\prime} \left(X _ {S}, X _ {i} \mid c, c ^ {\prime}\right) = L \left(X _ {S} ^ {\{c, c ^ {\prime} \}}, X _ {i} ^ {\{c, c ^ {\prime} \}}\right) \\ - \left(L \left(X _ {S} ^ {(c)}, X _ {i} ^ {(c)}\right) + L \left(X _ {S} ^ {(c)}, X _ {i} ^ {(c)}\right)\right). \tag {12} \\ \end{array}
$$
Given that the gains should not be positive by no-hypercompression, we minimize the pairwise scores  $L'$ . This can be seen as a graph clustering problem on a graph with nodes  $c \in C$ , weighted edges between each pair  $c, c'$  with weights  $L$  taking positive values iff  $c, c'$  belong to the same cluster, and where the number of clusters is unknown. Miyauchi et al. (2018) consider exactly this problem and show that it corresponds to the following Integer Linear Program (ILP) objective,
$$
\min  \sum_ {c _ {i}, c _ {j}} L ^ {\prime} \left(c _ {i}, c _ {j}\right) x _ {i j} \tag {13}
$$
where  $x_{ij}$  are decision variables indicating group membership,  $x_{ij} = 1$  if  $c_i, c_j \in \pi$  and zero otherwise. The minimization is subject to so-called triangle constraints [53] ensuring that the solution is a valid partition of  $\mathcal{C}$ , and can be solved using standard solvers.
As another variant, we can use hierarchical clustering with distance measure  $L'$ . With this approach, we discover the best partition of a given size  $k = |\Pi|$ . We select the best  $k$  by comparing the full MDL scores of the partitions at  $k$  and returning the minimal one.
We show how to adapt LINC in Alg. 2, and refer to it as  $\mathrm{LINC}_{\mathrm{CLUS}}$
B.4 LINC with Random Fourier Features
Last, we show how our score can be adapted to a common approximation of GPs using Random Fourier Features (RFFs) [34].
Computing  $L$  for GPs is computationally expensive since GP regression does not scale well with the number of samples. One way to mitigate the complexity of GP regression is to express its kernel through a Fourier transform, which in turn can be approximated efficiently. This is known as the Random Fourier Feature (RFF) approximation [34].
The approach is as follows: if  $\kappa$  is shift-invariant, we have
$$
\begin{array}{l} \kappa (x, y) = \int p (\omega) \cos (\omega^ {\top} (x - y)) d \omega \\ \approx \frac {1}{R} \sum_ {r = 1} ^ {R} \cos (i \omega_ {r} ^ {\top} (x - y)) = z (x) ^ {\top} z (y) \\ \end{array}
$$
for some probability measure  $p(\omega)$ . For the RBF kernel, this measure  $p(\omega)$  is another Gaussian distribution [18]. Here,  $z(x) = \frac{1}{\sqrt{R}} (\sqrt{2}\cos (\omega_1^\top x),\dots,\sqrt{2}\cos (\omega_R^\top x))^{\top}$  and  $z(y)$  similarly. Thus, for each pair  $x_{i},x_{j}$ , we obtain an unbiased estimator
$$
\kappa (x _ {i}, x _ {j}) = z (x _ {i}) ^ {\top} z (x _ {j}).
$$
The RFF approximation of GP prediction uses the regression  $f(x) = z(x)^{\top}w$ . We define our MDL encoding using the log-likelihood of this regression, as well as the norm
$$
\left\| f \right\| _ {\kappa} ^ {2} = w ^ {\top} Z w
$$
for the feature matrix  $Z = z(X)z(X)^{\top}$ . Up to these modifications, we use the MDL encoding in Eq. 6 for RFF regression models.
B.5 Complexity
For DAG discovery, exact search is infeasible due to the super-exponential search space over causal DAGs [33]. There are greedy solutions in  $\mathcal{O}(|\mathcal{V}|^3)$ ,  $\mathcal{V}$  being the number of nodes in the DAG [26].
Regarding the partition discovery problem, the exact solution involves the evaluation of  $b$  partitions,  $b$  being a bell number in  $|\mathcal{C}|$ . Heuristic search requires only  $|\mathcal{C}|^2$  regressions. While the problem of finding an optimal clustering of  $\mathcal{C}$  (e.g., using the ILP) is still NP-hard [53], there exist efficient options such as  $k$ -means that we could insert. In addition, we have  $|\mathcal{C}| \ll N$ , and ILP search can be done using standard solvers, which we found to be efficient in practice.