SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shuai Li, Yingjie Zhang
KLATASDS-MOE
Hongtu Zhu
Departments of Biostatistics,
Statistics, Computer Science, and Genetics
The University of North Carolina at Chapel Hill
Christina Dan Wang
Business Division
New York University Shanghai
Hai Shu
Department of Biostatistics
School of Global Public Health
Ziqi Chen* Zhuoran Sun, Yanfeng Yang
Conditional independence (CI) testing is a fundamental task in statistics and machine learning, but its effectiveness is hindered by the challenges posed by high-dimensional conditioning variables and limited data samples. This article introduces a novel testing approach to address these challenges and enhance control of the type I error while achieving high power under alternative hypotheses. The proposed approach incorporates a computationally efficient classifier-based conditional mutual information (CMI) estimator, capable of capturing intricate dependence structures among variables. To approximate a distribution encoding the null hypothesis, a  $k$ -nearest-neighbor local sampling strategy is employed. An important advantage of this approach is its ability to operate without assumptions about distribution forms or feature dependencies. Furthermore, it eliminates the need to derive asymptotic null distributions for the estimated CMI and avoids dataset splitting, making it particularly suitable for small datasets. The method presented in this article demonstrates asymptotic control of the type I error and consistency against all alternative hypotheses. Extensive analyses using both synthetic and real data highlight the computational efficiency of the proposed test. Moreover, it outperforms existing state-of-the-art methods in terms of type I and II errors, even in scenarios with high-dimensional conditioning sets. Additionally, the proposed approach exhibits robustness in the presence of heavy-tailed data.
Testing for conditional independence (CI) is a crucial and challenging task in statistics and machine learning, with wide-ranging applications in graphical models [26, 17], causal inference [44, 37, 22], and variable selection [11, 24]. The objective is to determine whether two random variables,  $X$  and  $Y$ , are independent given a set of conditioning variables  $Z$ , based on observations of the joint
*Corresponding author (zqchen@fem.ecnu.edu.cn).
distribution  $p_{X,Y,Z}(x,y,z)$ . Specifically, the hypothesis to be tested is:
$$
H _ {0}: X \perp Y | Z \text {v e r s u s} H _ {1}: X \not \perp Y | Z. \tag {1}
$$
However, CI testing becomes challenging due to the high-dimensionality of the conditioning variables  $Z$ , and the limited availability of data samples [5, 35, 40, 1]. Existing tests may struggle to control the type I error, particularly when handling high-dimensional conditioning variable sets with complex dependency structures [6, 40]. Moreover, even when a test is valid, the limited data availability can make it exceedingly challenging to distinguish between null and alternative hypotheses, resulting in low testing power [40].
In this article, we present a novel conditional independence (CI) testing method based on a  $k$ -nearest-neighbor local sampling scheme. Our approach incorporates two essential components: a conditional mutual information (CMI) estimator utilizing classification techniques, and a  $k$ -nearest-neighbor local sampling strategy to approximate the conditional distribution  $p_{X|Z}(x|z)$ . This approximation enables us to simulate data sets from a distribution that represents the null hypothesis, allowing us to estimate the CMI using these simulated data sets. By comparing the CMI estimator based on real data with those computed from the simulated data sets, we make informed decisions regarding the hypothesis testing. Theoretical analysis demonstrates that our proposed test achieves a valid control of the type I error asymptotically and exhibits consistency against all alternatives in  $H_1$ . Synthetic and real data analyses showcase that our test outperforms previous methods in terms of both the type I error and power under  $H_1$ , even when dealing with high-dimensional conditioning sets and/or small datasets. Furthermore, our approach remains computationally efficient as the dimension of the conditioning set and/or sample size increase. Additionally, our method is robust even in the presence of heavy-tailed data.
2.1 CMI estimation
In our work, we utilize CMI as a measure of conditional dependence. CMI offers a strong theoretical guarantee for CI testing, where  $I(X;Y|Z) = 0 \iff X \perp Y|Z$  [14]. It has the ability to capture complex dependence structures among variables, even in non-linear scenarios. Estimators based on  $k$ -nearest neighbors and kernel methods have been widely employed for CMI estimation [27, 38, 23, 19]. However, these methods may encounter efficiency issues when dealing with high-dimensional conditioning variables, known as the curse of dimensionality [34, 32]. [32] approached the CMI estimation problem by formulating it as a minimax optimization problem and proposed the use of generative adversarial networks (GANs) to optimize and obtain the CMI estimator, which can handle high-dimensional scenarios. However, the training of GANs is often challenging, with the risk of collapse if hyperparameters and regularizers are not carefully chosen [15]. Recently, [34, 30] proposed a classifier-based CMI estimator capable of handling high dimensions. They initially developed an estimator for the Kullback-Leibler (KL) divergence using a classifier and then derived mutual information (MI) estimators based on it. CMI is defined as the difference between two MI values, and the CMI estimator is obtained by computing the difference between two MI estimators. In contrast, our method directly utilizes the classifier-based KL-divergence estimator to obtain the CMI estimator. This approach also handles high dimensions effectively and is computationally more efficient compared to the method employed by [34, 30].
2.2 CI Testing
In recent years, a considerable body of literature on CI testing has emerged. Here, we provide a brief overview of some existing methods, and for a more comprehensive review, we refer readers to [28].
One important category of CI testing methods involves using different measures for CI [3, 46, 2]. For example, [18] proposed a CI testing method based on the empirical Hilbert-Schmidt norm of the conditional cross-covariance operator. Another approach, KCIT, was introduced by [47], which utilizes the partial association of reproducing kernel Hilbert spaces to measure conditional independence. [40] proposed a novel kernel-based CI testing method using the  $l_{p}$  distance between two well-chosen analytic kernel mean embeddings evaluated at a finite set of locations. [42] proposed the generalized covariance measure for CI testing based on a regression method. However, obtaining the exact distribution of the test statistic derived from the conditional independence measure under
$H_0$  can be challenging. Instead, researchers have employed alternative methods to derive asymptotic distributions of the test statistics under the null hypothesis [47, 46, 45, 40]. Nevertheless, the effectiveness of asymptotic distributions can be compromised when the sample size is small or the dimension of  $Z$  is high [16, 38]. As a result, tests based on asymptotic distributions may exhibit inflated type-I errors or inadequate power when dealing with small sample sizes or high-dimensional  $Z$  [47, 40]. In our work, we choose to utilize CMI as a metric for conditional dependence.
To overcome the reliance on asymptotic null distributions, [11] proposed the model-X framework, assuming knowledge of the conditional distribution  $p_{X|Z}(x|z)$ . Under this assumption, a set of test statistics can be computed that are exchangeable under the null hypothesis, either through direct resampling [11] or permutation methods [8]. However, in practice, the distribution of  $X|Z$  is rarely available. Therefore, accurately approximating the distribution of  $X|Z$  becomes crucial for maintaining the type I error control of these tests. To address this challenge, [5] developed the generative conditional independence test (GCIT) using Wasserstein GANs to approximate the distribution of  $X|Z$ , while [43] proposed using Sinkhorn GANs for the same purpose. However, as highlighted in [5, 30], limited data and noise may lead to inaccurate learning of conditional distributions using neural networks, resulting in inflated type I errors for GCIT and the double GANs-based CI test (DGCIT). To mitigate this issue, [30] proposed using the 1-nearest neighbor method to generate samples from the approximated conditional distribution of  $X$  given  $Z$ . However, their method requires splitting the dataset into two parts, with the testing dataset used for calculating the test statistics comprising only one-third of the total samples. This will reduce the statistical power of the test, particularly when working with small datasets. [38] utilized a Kozachenko-Leonenko estimator for CMI as the test statistic and proposed a permutation scheme based on  $k$ -nearest neighbors to generate samples under the null distribution. However, the curse of dimensionality adversely affects the performance of the Kozachenko-Leonenko CMI estimator, leading to poor performance when the conditioning variable set  $Z$  has high dimensions. Additionally, no theoretical guarantee is provided. [31] and [25] employed a binning strategy, discretizing  $Z$  into a finite number of bins based on the proximity of conditioning variables  $Z$ , followed by the "permute-within-groups" strategy. However, selecting bins in high-dimensional settings presents a significant challenge [8], and their methods are limited to handling conditioning variables  $Z$  with very few dimensions.
In our work, we propose the utilization of a  $k$ -nearest-neighbor local sampling strategy as an appealing alternative to the binning strategy. Specifically, our strategy generates samples locally based on the  $k$ -nearest neighbors of the conditioning variables  $Z$ . We demonstrate that the distribution of samples generated from this  $k$ -nearest-neighbor local sampling scheme closely approximates the true conditional distribution  $p_{X|Z}(x|z)$  in terms of total variation distance. One significant advantage of our proposed method is that it allows the entire dataset to be used for computing the testing statistics, eliminating the need for dataset splitting. This feature makes our method more effective compared to [30], which requires dataset splitting, particularly when dealing with small datasets. Moreover, our method does not require the derivation of asymptotic null distributions for the estimated CMI, and it can easily handle high-dimensional conditioning variables  $Z$ . We provide theoretical and empirical evidence that our test achieves a valid control of the type I error and attains high power under the alternative hypothesis  $H_1$ .
3  $K$ -Nearest-Neighbor Local Sampling Based CI Testing
3.1 Classifier-based CMI estimator
The CMI for a triplet of random variables/vectors  $(X,Y,Z)$  is defined as:
$$
I (X; Y | Z) = \iiint p _ {X, Y, Z} (x, y, z) \log \frac {p _ {X , Y , Z} (x , y , z)}{p _ {X , Z} (x , z) p _ {Y | Z} (y | z)} d x d y d z, \tag {2}
$$
where  $p_{X,Y,Z}(x,y,z)$  is the joint density of  $(X,Y,Z)$ ,  $p_{X,Z}(x,z)$  is the joint density of  $(X,Z)$ , and  $p_{Y|Z}(y|z)$  is the conditional density of  $Y$  given  $Z = z$ . One approach to estimate CMI is by directly estimating the joint and conditional densities from the available data and plugging them into (2). However, accurately estimating the density functions can be challenging, especially in high-dimensional settings where it is more difficult than directly estimating CMI (2) [34]. CMI is a special case of the Kullback-Leibler (KL) divergence, and thus we have:
$$
I (X; Y | Z) = D _ {K L} \left(p _ {X, Y, Z} (x, y, z) \mid \mid p _ {X, Z} (x, z) p _ {Y | Z} (y | z)\right), \tag {3}
$$
Algorithm 1 1-Nearest-Neighbor sampling  $(\mathbf{1} - \mathbf{NN}(V_1,V_2,n))$
Input: Datasets  $V_{1}$  and  $V_{2}$ , both with sample size  $n$  and  $V = V_{1} \cup V_{2}$  consisting of  $2n$  independently and identically distributed (i.i.d.) samples from  $p_{X,Y,Z}(x,y,z)$ .
Output: A new data set  $V'$  consists of  $n$  samples.
1: Let  $V' = \emptyset$ .
2: for  $(X,Y,Z)$  in  $V_{2}$  do
3: Go to  $V_{1}$  to find the sample  $(X', Y', Z')$  such that  $Z'$  is the 1-nearest neighbor of  $Z$  in terms of the  $l_{2}$  norm
4:  $V^{\prime} = V^{\prime}\cup \{(X,Y^{\prime},Z)\}$
5: end for
6: return  $V^{\prime}$
where  $D_{KL}(f||g)$  denotes the KL divergence between two distribution functions  $F$  and  $G$ , with density functions  $f(x)$  and  $g(x)$ , respectively. The Donsker-Varadhan (DV) representation of  $D_{KL}(f||g)$  is given by:
$$
\sup  _ {s \in S} \left[ E _ {w \sim f} s (w) - \log E _ {w \sim g} \exp \{s (w) \} \right], \tag {4}
$$
where the function class  $\mathcal{S}$  includes all functions with finite expectations. The optimal function in (4) is given by  $s^* (x) = \log \{f(x) / g(x)\}$  [4], which leads to:
$$
D _ {K L} (f | | g) = E _ {w \sim f} \log \left\{f (w) / g (w) \right\} - \log \left[ E _ {w \sim g} \left\{f (w) / g (w) \right\} \right]. \tag {5}
$$
Building upon the classifier-based CMI estimation method proposed by [34], we propose a CMI estimation method using the 1-nearest-neighbor (1-NN) sampling algorithm [41, 30]. The pseudocode for estimating CMI is outlined in Algorithm 2. The main objective is to empirically estimate (5) with  $f = p_{X,Y,Z}(x,y,z)$  and  $g = p_{X,Z}(x,z)p_{Y|Z}(y|z)$ , which requires samples from both  $p_{X,Y,Z}(x,y,z)$  and  $p_{X,Z}(x,z)p_{Y|Z}(y|z)$ . The available data only consists of samples from  $p_{X,Y,Z}(x,y,z)$ . However, generating samples from  $p_{X,Z}(x,z)p_{Y|Z}(y|z)$  requires knowledge of the unknown conditional distribution  $p_{Y|Z}(y|z)$ . To address this challenge, we propose utilizing the 1-NN sampling algorithm [41, 30], which is outlined in Algorithm 1. Empirical and theoretical results presented in [30] indicate that the 1-NN sampling algorithm can accurately approximate the conditional distribution.
Next, we formalize the classifier-based CMI estimator. We consider a data set  $V$  consisting of  $2n$  i.i.d. samples  $\{W_i := (X_i, Y_i, Z_i)\}_{i=1}^{2n}$  with  $(X_i, Y_i, Z_i) \sim p_{X,Y,Z}(x,y,z)$ . The data set  $V$  is divided into two equally sized parts  $V_1$  and  $V_2$ , where  $|V_1| = |V_2| = n$ . For data sets  $V_1$  and  $V_2$ , we use the 1-NN sampling algorithm 1 to generate a new data set  $V'$  with  $n$  samples. We assign labels  $l = 1$  for all samples in  $V_2$  and  $l = 0$  for all samples in  $V'$ . In this supervised classification task, a binary classifier can be trained using an advanced binary classification model, such as XGBoost [41, 12] or deep neural networks [21]. The classifier produces predicted probability  $\alpha_m = P(l = 1|W_m)$  for a given sample  $W_m$ , leading to an estimator of the likelihood ratio on  $W_m$  given by  $\widehat{L}(W_m) = \alpha_m / (1 - \alpha_m)$ . It follows from (3) and (5) that an estimator of  $I(X;Y|Z)$  is given by
$$
\begin{array}{l} \widehat {I} (X; Y | Z) := \widehat {D} _ {K L} (p _ {X, Y, Z} (x, y, z) | | p _ {X, Z} (x, z) p _ {Y | Z} (y | z)) \\ = d ^ {- 1} \sum_ {i = 1} ^ {d} \log \widehat {L} \left(W _ {i} ^ {f}\right) - \log \left\{d ^ {- 1} \sum_ {j = 1} ^ {d} \widehat {L} \left(W _ {j} ^ {g}\right) \right\}, \tag {6} \\ \end{array}
$$
where  $d = \lfloor n / 3\rfloor$  with  $\lfloor t\rfloor$  being the largest integer not greater than  $t$ ,  $W_{i}^{f}$  is a sample in  $V_{f}^{\mathrm{test}}$  and  $W_{j}^{g}$  is a sample in  $V_{g}^{\mathrm{test}}$ , where  $V_{f}^{\mathrm{test}}$  and  $V_{g}^{\mathrm{test}}$  are defined in Algorithm 2. According to Theorem 1 in [34],  $\widehat{I} (X;Y|Z)$  is a consistent estimator of  $I(X;Y|Z)$ .
In contrast to the approach taken by [34, 30], which estimate CMI by using the difference between two mutual informations (i.e.,  $I(X;Y|Z) = I(X;Y,Z) - I(X;Z)$ ) and therefore require training two binary classifiers, our method achieves CMI estimation with just a single binary classifier. This leads to significantly improved computational efficiency. As depicted in Algorithm 3, calculating a single  $p$ -value requires the computation of  $(B + 1)$  CMIs. Hence, in practical applications involving real data analysis, especially when dealing with large sample sizes or high-dimensional conditional variables, efficient computation of the CMI becomes crucial.
Algorithm 2 Classifier-based CMI Estimator
Input: Dataset  $V$  containing  $2n$  i.i.d. samples drawn from  $p_{X,Y,Z}(x,y,z)$ .
Output: An estimator of CMI.
1: Equally split  $V$  into two parts  $V_{1}$  and  $V_{2}$ , each containing  $n$  samples.
2: Apply Algorithm 1 to generate a new dataset  $V'$  with  $|V'| = n$ .
3: Form the labeled datasets  $V_{f} = \{(W_{i}^{f},l = 1):W_{i}^{f}\in V_{2}\}$  and  $V_{g} = \{(W_{j}^{g},l = 0):W_{j}^{g}\in V^{\prime}\}$ .
4: Divide  $V_{f}$  into training and testing subsets  $V_{f}^{\text{train}}$  and  $V_{f}^{\text{test}}$ , at a ratio of 2:1.
5: Similarly, split  $V_{g}$  into training and testing subsets  $V_{g}^{\text{train}}$  and  $V_{g}^{\text{test}}$ , at a ratio of 2:1.
6: Merge the datasets to form  $V^{\text{train}} = V_f^{\text{train}} \cup V_g^{\text{train}}$  and  $V^{\text{test}} = V_f^{\text{test}} \cup V_g^{\text{test}}$ .
7: Train the classifier  $C$  using  $V^{\mathrm{train}}$ .
8: For each  $w \in V_0^{test}$ , where  $V_0^{test}$  includes all features in  $V^{test}$ , compute the classifier-based predicted probability  $P(l = 1|w)$ .
9: Calculate  $\widehat{I}(X;Y|Z)$  as per formula (6).
10: return  $\widehat{I}(X;Y|Z)$ .
3.2 The Proposed CI Testing Procedure
While we have proposed a consistent estimator for CMI, accurately estimating it as zero under the null hypothesis,  $H_0$ , is practically unattainable due to random errors in the sample data. These errors lead to deviations of the estimator from the actual value. To enhance the effectiveness of CI testing based on the CMI estimator, we propose a test that takes into account the statistical variation inherent in the CMI estimator.
Under  $H_0$ , we can express the following representation:
$$
p _ {X \mid Y, Z} (x \mid y, z) = p _ {X \mid Z} (x \mid z). \tag {7}
$$
The null model in (7) preserves the dependence between  $X$  and  $Z$ , while breaking any dependence between  $X$  and  $Y$ . Therefore, if a direct causal link exists between  $X$  and  $Y$ , replacing  $X$  with a null sample  $X^{*} \sim p_{X|Z}(x|z)$  would disrupt this relationship. Thus, we can conclude that  $X^{*}$  and  $Y$  are conditionally independent given  $Z$ , i.e.,  $I(X^{*};Y|Z) = 0$ .
Consider  $n$  i.i.d. copies  $\mathcal{A} = \{(x_i, y_i, z_i) : i = 1, \dots, n\}$  of  $(X, Y, Z)$ . When the distribution of  $X|Z$  is known, conditional on  $Z = (z_1, \ldots, z_n)^T$ , we can independently draw a pseudo sample  $x_i^{(b)} \sim p_{X|Z}(x|z_i)$  for each  $i$  across  $b = 1, \dots, B$  such that all  $X^{(b)} := (x_1^{(b)}, \ldots, x_n^{(b)})^T$  are independent of  $Y := (y_1, \ldots, y_n)^T$  and also  $X := (x_1, \ldots, x_n)^T$ , where  $B$  is the number of repetitions. Under  $H_0$ ,  $(X^{(b)}, Y, Z) \stackrel{d}{=} (X, Y, Z)$  for all  $b$ , where  $d$  denotes equality in distribution. We denote the CMI estimator of  $I(X; Y|Z)$  based on  $(X^{(b)}, Y, Z)$  as  $\widehat{\mathrm{CMI}}^{(b)}$  and denote the estimator based on  $(X, Y, Z)$  as  $\widehat{\mathrm{CMI}}$ . We can approximate the  $p$ -value by
$$
p := \frac {1 + \sum_ {b = 1} ^ {B} 1 (\widetilde {\mathrm {C M I}} ^ {(b)} \geq \widehat {\mathrm {C M I}})}{1 + B}, \tag {8}
$$
where  $1(\cdot)$  is the indicator function. To begin with, we demonstrate that the test based on (8) can effectively control the type I error. Specifically, under the null hypothesis  $H_0$ , the  $(B + 1)$  triples  $(X,Y,Z)$ ,  $(X^{(1)},Y,Z),\ldots ,(X^{(B)},Y,Z)$  are exchangeable, and thus the  $p$ -value is valid and satisfies  $P(p\leq \alpha |H_0)\leq \alpha$  for any given  $\alpha \in (0,1)$  [11, 8, 43]. Furthermore, it intuitively suggests that our test can achieve high power under the alternative hypothesis  $H_{1}$ . As  $n$  approaches infinity,  $\widetilde{\mathrm{CMI}}^{(b)}$  converges to zero in probability. But under  $H_{1}$ , we know that  $I(X;Y|Z) > 0$ . As a result,  $\widehat{\mathrm{CMI}}$  should be positive with high probability. Consequently, the  $p$ -value calculated using (8) is very small with high probability, indicating the consistency of our test against all alternatives stated in  $H_{1}$ .
However,  $p_{X|Z}(x|z)$  is rarely known in practice [5, 43]. We propose using the  $k$ -nearest-neighbor local sampling strategy to approximate  $p_{X|Z}(x|z)$ . The approximated distribution is denoted as  $\widehat{p}_{X|Z}(x|z)$ . To generate pseudo samples from  $\widehat{p}_{X|Z}(x|z)$ , we follow the steps below:
Algorithm 3  $K$ -Nearest-neighbor local sampling based CI testing
Input: Dataset  $(X,Y,Z)$  consisting of  $n$  i.i.d. samples from  $p_{X,Y,Z}(x,y,z)$ .  
Parameter: The number of repetitions  $B$ ; the neighbor order  $k$ ; the significance level  $\alpha$ .  
Output: Accept  $H_0:X\perp Y|Z$  or  $H_{1}:X\not\perp Y|Z$ .  
1: Use Algorithm 2 to obtain  $\widehat{\mathbf{CMI}}$  based on  $(X,Y,Z)$ .  
2: for  $i\in \{1,2,\dots,n\}$  do  
3: Obtain the set of indices of the  $k$ -nearest neighbor of  $z_{i}$ :  $\mathcal{M}_i = \{j\in \{1,2,\dots,n\} : \| z_j - z_i\| _2\leq d_i^k\}$ , where  $d_i^k$  denotes the distance of  $z_{i}$  to its  $k$ -nearest neighbor.  
4: end for  
5:  $b = 1$ .  
6: while  $b\leq B$  do  
7: Initialize empty array  $\widetilde{\boldsymbol{X}}$  of length  $n$ .  
8: for  $i\in \{1,2,\dots,n\}$  do  
9: Shuffle  $\mathcal{M}_i$   
10:  $j = \mathcal{M}_i[1]$   
11:  $\widetilde{\boldsymbol{X}} [i] = \boldsymbol {X}[j]$ .  
12: end for  
13: Use Algorithm 2 to obtain  $\widehat{\mathbf{CMI}}^{(b)}$  based on  $(\widetilde{X},Y,Z)$ .  
14:  $b = b + 1$ .  
15: end while  
16: Compute  $p$ -value:  $p\coloneqq [1 + \sum_{b = 1}^{B}1\{\widehat{\mathbf{CMI}}^{(b)}\geq \widehat{\mathbf{CMI}}\} ] / (1 + B)$ .  
17: if  $p\geq \alpha$  then  
18: Accept  $H_0:X\perp Y|Z$ .  
19: else  
20: Accept  $H_{1}:X\not\perp Y|Z$ .  
21: end if
1. Obtain the set of indices of the  $k$ -nearest neighbors of  $z_{i}$  based on the  $l_{2}$  norm. Denote this set as  $\mathcal{M}_i = \{j\in \{1,\dots,n\} : \| z_j - z_i\| _2\leq d_i^k\}$ , where  $d_i^k$  is the distance of  $z_{i}$  to its  $k$ -nearest neighbor;
2. Shuffle  $\mathcal{M}_i$  and let  $j$  be the first element of  $\mathcal{M}_i$ . Then, define the pseudo sample  $\tilde{x}_i$  as  $x_j$ ;
3. Repeat the above two steps for all  $i \in \{1, 2, \ldots, n\}$  to obtain the resultant pseudo sample vector  $\widetilde{\boldsymbol{X}} = (\tilde{x}_1, \dots, \tilde{x}_n)^T$ .