Approximate Kernel-Based Conditional Independence Tests for Fast Non-Parametric Causal Discovery

1 The problem

Constraint-based causal discovery (CCD) algorithms such as Peter-Clark (PC) and Fast Causal Inference (FCI) infer causal relations from observational data by combining the results of many conditional independence (CI) tests [1]. In practice, a CCD algorithm can easily request p-values from thousands of CI tests even with a sparse underlying graph. Developing fast and accurate CI tests is therefore critical for maximizing the usability of CCD algorithms across a wide variety of datasets.

Investigators have developed many fast parametric methods for testing CI. For example, we can use partial correlation to test for CI under the assumption of Gaussian variables [2], [3]. We can also consider testing for unconditional independence X⫫Y|Z=z for each constant value z when Z is discrete and P(Z=z)>0. The chi-squared test for instance utilizes this strategy when both X and Y are also discrete [4]. Another permutation-based test generalizes the same strategy even when X and Y are not necessarily discrete [5].

Testing for CI in the non-parametric setting generally demands a more sophisticated approach. One strategy involves discretizing continuous conditioning variables Z as Z˘ in some optimal fashion and assessing unconditional independence ∀Z˘=z˘ [6], [7]. Discretization however suffers severely from the curse of dimensionality because consistency arguments demand smaller bins with increasing sample size, but the number of cells in the associated contingency table increases exponentially with the conditioning set size. A second method involves measuring the distance between estimates of the conditional densities f(X|Y,Z) and f(X|Z), or their associated characteristic functions, by observing that f(X|Y,Z)=f(X|Z) when X⫫Y|Z [8], [9]. However, the power of these tests also deteriorates quickly with increases in the dimensionality of Z.

Several investigators have since proposed reproducing kernel-based CI tests in order to tame the curse of dimensionality. Indeed, kernel-based methods in general are known for their strong empirical performance in the high dimensional setting. The Kernel Conditional Independence Test (KCIT) for example assesses CI by capitalizing on a characterization of CI in reproducing kernel Hilbert spaces (RKHSs; [10]). Intuitively, KCIT works by testing for vanishing regression residuals among functions in RKHSs. Another kernel-based CI test called the Permutation Conditional Independence Test (PCIT) reduces CI testing to two-sample kernel-based testing via a carefully chosen permutation found at the solution of a convex optimization problem [11].

The aforementioned kernel-based CI tests unfortunately suffer from an important drawback: both tests scale at least quadratically with sample size and therefore take too long to return a p-value in the large sample size setting. In particular, KCIT’s bottleneck lies in the eigendecomposition as well as the inversion of large kernel matrices [10], and PCIT must solve a linear program that scales cubicly with sample size in order to obtain its required permutation [11]. As a general rule, it is difficult to develop exact kernel-based methods which scale sub-quadratically with sample size, since the computation of kernel matrices themselves scales at least quadratically.

Many investigators have nonetheless utilized random Fourier features in order to quickly approximate kernel methods in linear time with respect to the number of Fourier features. For example, Lopez-Paz and colleagues developed an unconditional independence test using statistics obtained from canonical correlation analysis with random Fourier features [12]. Zhang and colleagues have also utilized random Fourier features for unconditional independence testing but took a different approach by approximating the kernel cross-covariance operator [13]. The authors further analyzed block and Nyström-based kernel approximations to the unconditional independence testing problem. The authors ultimately concluded that the random Fourier feature and the Nyström-based approaches both outperformed the block-based approach on average. Others have analyzed the use of random Fourier features for predictive modeling (e. g., [14], [15]) or dimensionality reduction [16]. In practice, investigators have observed that methods which utilize random Fourier features often scale linearly with sample size and achieve comparable accuracy to exact kernel methods [12], [13], [14], [15], [16].

In this paper, we also use random Fourier features to design two fast tests called the Randomized Conditional Independence Test (RCIT) and the Randomized conditional Correlation Test (RCoT) which approximate the solution of KCIT. Simulations show that RCIT, RCoT and KCIT have comparable accuracy, but both RCIT and RCoT scale linearly with sample size in practice. As a result, RCIT and RCoT return p-values several orders of magnitude faster than KCIT in the large sample size context. Moreover, experiments demonstrate that the causal structures returned by CCD algorithms using either RCIT, RCoT or KCIT have nearly identical accuracy.

2 High-level summary

We now provide a high-level summary for investigators who simply wish to perform CCD with RCIT or RCoT. We want to test whether we have X⫫Y|Z in a fast and accurate manner without resorting to parametric assumptions. Previously, Zhang and colleagues introduced a non-parametric CI test called KCIT which can test whether we have X⫫Y|Z in an accurate but not fast manner by analyzing the partial cross-covariance (operator) using the kernel method [10]. Kernels however are expensive to compute because they scale at least quadratically with sample size. Another line of work has fortunately shown that we can approximate kernels by averaging over a variety of non-linear transformations called random Fourier features (e. g., [14], [15], [17]). We therefore propose to approximate KCIT by utilizing random Fourier features, specifically by analyzing the partial cross-covariance matrix of {X,Z} and Y (RCIT) or X and Y (RCoT) after subjecting the variable sets to the non-linear transformations and then non-linearly regressing out the effect of Z. Simulations show that RCIT and RCoT return p-values in a much shorter time frame while also matching or outperforming KCIT in approximating the null distribution. RCoT in particular also returns the most accurate p-values when the conditioning set size Z is large (≥4).

3 Preliminaries on kernels

We will deal with kernels in this paper, so we briefly review the corresponding theory; see [18] for a more extensive discussion of similar concepts. Capital letters X,Y,Z denote sets of random variables with codomains X,Y,Z, respectively. Let HX correspond to a Hilbert space of functions mapping X to R. We say that HX is more specifically a reproducing kernel Hilbert space, if the Dirac delta operator δx:HX↦R (which maps f∈HX to f(x)∈R) is a bounded linear functional [19]. We can associate HX with a unique positive definite kernelkX:X×X↦R in which the feature mapϕ(x):X↦HX satisfies the reproducing property⟨f,ϕ(x)⟩HX=f(x), ∀f∈HX, ∀x∈X by the Moore-Aronszajn Theorem. We will require that HX be separable (i. e., it must have a complete orthonormal system) throughout this paper [18]. Hein and Bousquet [20] showed that any continuous kernel on a separable space X (e. g., Rd) induces a separable RKHS. We will likewise consider other kernels such as kY on the separable space Y.

We have the following norm defined on linear operators between RKHSs:

We denote the probability distribution of X as PX and that of Y as PY. We may then define the mean elementsE[ϕ(X)]=μX∈HX and E[ψ(Y)]=μY∈HY with respect to the probability distributions; here, ϕ denotes the feature map from X to HX, and ψ likewise denotes the feature map from Y to HY. We also define the quantity ‖μX‖HX2 by applying the expectation twice:

where X′ is an independent copy of X which follows the same distribution. The mean elements μX and μY exist so long as their respective norms in HX or HY are bounded; this is true so long as the kernels kX and kY are also bounded so that EXX′[kX(X,X′)]<∞ and EYY′[kY(Y,Y′)]<∞ [18].

The cross-covariance operator associated with the joint probability distribution PXY over (X,Y) is a linear operator ΣXY:HY↦HX defined as follows:

for all f∈HX and g∈HY. Gretton et al. [18] showed that the HS-norm of the cross covariance operator, ‖ΣXY‖HS2, can be written in terms of kernels as follows:

provided that X⫫Y, and kX as well as kY are centered (i. e., EXX′[kX(X,X′)]=0 and EYY′[kY(Y,Y′)]=0). We can therefore consider the following empirical estimate of the ‖ΣXY‖HS2 using n i. i. d. samples x,y as follows, if we assume that X⫫Y [10], [18]:

where Txy→p‖ΣXY‖HS2 [18]; the dependent case can be found in Lemma 1 of the same paper for the interested reader. The notations K˜X and K˜Y correspond to centralized kernel matrices such that:

and likewise for K˜Y. Here, H=I−1n11T, I denotes an n×n identity matrix, and 1 denotes a vector of ones. The notation KX denotes a kernel matrix such as the RBF kernel where [KX]ij=exp(−‖xi−xj‖22σ) with xi,xj∈x. The transformation in Equation 6 ensures that 1n2∑i=1n∑j=1n[K˜X]ij=0 similar to the centered kernels kX and kY in Equation 4.

4 Characterizations of conditional independence

We denote the probability distribution of X as PX and the joint probability distribution of (X,Z) as PXZ. Let LX2 denote the space of square integrable functions of X, and LXZ2 that of (X,Z). Here, LX2={s(X)∣EX(|s|2)<∞} and likewise for LXZ2. Next consider a dataset of n i. i. d. samples drawn according to PXYZ.

We use the notation X⫫Y|Z when X and Y are conditionally independent given Z. Perhaps the simplest characterization of CI reads as follows: X⫫Y|Z if and only if PXY|Z=PX|ZPY|Z. Equivalently, we have PX|YZ=PX|Z and PY|XZ=PY|Z.

5 Kernel conditional independence test

We now consider the following hypotheses:

We may equivalently rewrite the above null and alternative more explicitly using Proposition 1 as follows, provided that the kernels are chosen such that the premises of that proposition are satisfied (e. g., when the kernels are Gaussian or Laplacian [22]):

The above hypothesis implies that we can test for conditional independence by testing for uncorrelatedness between functions in reproducing kernel Hilbert spaces.

Zhang et al. [ 10] exploited the equivalence between 11 and 12 in the Kernel Conditional Independence Test (KCIT), which we now describe. Consider the functional spaces f∈HX¨, g∈HY, and h∈HZ. We can compute the corresponding centered kernel matrices K˜X¨, K˜Y and K˜Z from n i. i. d. samples x,y and z as in 6. We can then use these matrices to perform kernel ridge regression in order to estimate the function h∗∈HZ in 10 as follows: hˆ∗(z)=K˜Z(K˜Z+λI)−1f(x¨), where λ denotes the ridge regularization parameter and f∈HX¨. Consequently, we can estimate the residual function f˜ as f˜(x¨)=f(x¨)−hˆ∗(z)=RZf(x¨) where:

Next consider the eigenvalue decomposition K˜X¨=VX¨ΛX¨VX¨. Let ϕX¨=VX¨ΛX¨1/2. Then the empirical kernel map of HX¨|Z is given by ϕ˜X¨=RZϕX¨ because K˜X¨|Z=ϕ˜X¨ϕ˜X¨T. We may similarly write K˜Y|Z=ϕ˜Yϕ˜YT. We can thus write the centralized kernel matrices corresponding to f˜ and g˜ as follows:

We can use the above two centered kernel matrices to compute an empirical estimate of the HS norm of the partial cross covariance operator similar to how we computed the quantity TXY in Equation 5:

Note that Tx¨y·z denotes an empirical estimate of ‖ΣX¨Y·Z‖HS2. KCIT uses the statistic SK=nTx¨y·z in order to determine whether or not to reject H0; we multiply the empirical estimate of the HS norm by n in order to ensure convergence to a non-degenerate distribution.

6 Proposed test statistic & its asymptotic distribution

Observe that computing SK requires the inversion of large kernel matrices which scales strictly greater than quadratically with sample size; the exact complexity depends on the algorithm used to compute the matrix inverse (e. g., O(n2.376) if we use the Coppersmith-Winograd algorithm [25]). CCD algorithms run with KCIT therefore take too long to complete. In this report, we will propose an inexpensive hypothesis test that almost always rejects or fails to reject the null whenever conditional dependence or independence holds, respectively, and even outperforms 12 as illustrated in our experimental results in Section 7. We will use a strategy that has already been successfully adopted in the context of unconditional dependence testing in doing so [12], [13].

We will in particular also take advantage of the characterization of CI presented in Proposition 1. Recall that the Frobenius norm corresponds to the Hilbert-Schmidt norm in Euclidean space. We therefore consider the following hypotheses as approximations to those in 12 and 11:

where CA¨B·Z=E[(A¨i−E(A¨|Z))(Bi−E(B|Z))T] corresponds to the ordinary partial cross covariance matrix, where E(A¨|Z) and E(B|Z) may be non-linear functions of Z. We also have A¨=f′(X¨)≜{f1′(X¨),…,fm′(X¨)} with fj′(X¨)∈GX¨, ∀j. Similarly, B=h′(Y)≜{h1′(Y),…,hq′(Y)} with hk′(Y)∈GY, ∀k. The terms GX¨ and GY denote two spaces of functions, which we set to be the support of the process 2cos(WT·+B), W∼PW, B∼Uniform([0,2π]). In other words, we select m functions from GX¨ and q functions from GY. We select these specific spaces because we can use them to approximate continuous shift-invariant kernels. A kernel k is said to be shift-invariant if and only if, for any a∈Rp, we have k(x−a,y−a)=k(x,y), ∀(x,y)∈Rp×Rp; examples of shift-invariant kernels include the Gaussian kernel frequently used in KCIT or the Laplacian kernel. The following result allows us to perform the approximation using the proposed spaces:

The precise form of PW depends on the type of shift-invariant kernel one would like to approximate (see Figure 1 of Rahimi and Recht [14] for a list). Since investigators most frequently implement KCIT using the Gaussian kernel kσ(x,y)=exp(−‖x−y‖2/σ) with hyperparameter σ, we choose to approximate the Gaussian kernel by setting PW to a centered Gaussian with standard deviation σ/2.

We will use the squared Frobenius norm of the empirical partial cross-covariance matrix as the statistic for RCIT:

where CˆA¨B·Z=1n−1∑i=1n[(A¨i−Eˆ(A¨|Z))(Bi−Eˆ(B|Z))T]. Recall however that E(A¨|Z) and E(B|Z) may be non-linear functions of Z and therefore difficult to estimate. We would instead like to approximate E(A¨|Z) and E(B|Z) with linear functions. We therefore let C=g(Z)≜{g1(Z),…,gd(Z)} with gl(Z)∈GZ,∀l, where we also set GZ to be the support of the process 2cos(WT·+B), W∼PW, B∼Uniform([0,2π]). We will approximate CA¨B·Z with CˆA¨B·C=CˆA¨B−CˆA¨C(ΣˆCC+γI)−1CˆCB similar to 7, where γ denotes a small ridge parameter; recall that this is equivalent to computing the cross-covariance matrix across the residuals of A¨ and B given C using linear ridge regression. We can justify this procedure because the particular choice of C allows us to approximate both E(A¨|Z) and E(B|Z) with linear functions of C as described below.

Let fj=fj′−E(fj′|Z). Then E(fj|Z)=0, so fj∈FXZ. Moreover, hk′−E(hk′|Z)∈FY·Z. Note that we can estimate E(fj′|Z) with the linear ridge regression solution uˆjTg(Z) under mild conditions because we can guarantee the following:

The exponential rate with respect to d in the above proposition suggests we can approximate the output of kernel ridge regression with a small number of random Fourier features, a hypothesis which we verify empirically in Section 7. Moreover, we can estimate E(hk′|Z) with uˆkTg(Z), because we can similarly guarantee that P[|E˘(hk′|Z)−uˆkTg(Z)|≥ε]→0 for any fixed ε>0 at an exponential rate with respect to d.

We can therefore consider the following spaces for S which are similar to the L2 spaces used in claim 4 of Proposition 2:

We then approximate CI with S in the following sense:

1. We always have X⫫Y|Z⟹E(fh)=0, ∀f∈GˆX¨and∀h∈GˆY·Z.

2. The reverse direction may hold for a larger subset of all possible joint distributions as the values of m and q increase; this is because at least one entry of CA¨B·C will likely be greater than zero for any given distribution as the values of m,q increase.