Toward Scalable and Valid Conditional Independence Testing with Spectral Representations

Technical Summary: Toward Scalable and Valid Conditional Independence Testing with Spectral Representations

Problem Statement
Conditional independence (CI) testing, the task of determining whether random variables $X$ and $Y$ are independent given a conditioning set $Z$ ($H_0: X \perp\!\!\perp Y | Z$), is fundamental to causal inference, feature selection, and graphical modeling. However, in nonparametric settings, CI testing is known to be fundamentally challenging. Shah & Peters (2020) established that no test can uniformly control Type I error across all conditionally independent distributions while maintaining power against all alternatives. Consequently, existing methods rely on restrictive structural assumptions (e.g., smoothness, known conditional distributions, or specific regression rates) to achieve validity.

While kernel-based methods using partial covariance operators (e.g., KCIT, RCIT) offer a principled approach by implicitly encoding structural assumptions, they suffer from limited adaptivity and poor scalability in high-dimensional settings. Conversely, methods like Model-X or local permutation tests require strong assumptions about the conditional distribution $P(X|Z)$ or smoothness of the data, which may not hold in practice. There is a need for a method that bridges the theoretical rigor of kernel methods with the scalability and adaptivity of modern representation learning.

Methodology
The authors propose SpectralCIT, a framework that leverages spectral representation learning to address the limitations of classical kernel tests. The core idea is to learn the leading spectral features (singular functions) of the partial covariance operator $\Sigma_{X\dot{Y} \cdot Z}$, where $\dot{Y} = (Y, Z)$, and use these learned representations to construct a simple, scalable test statistic.

1. Spectral Representation Learning:
   The method aims to learn the rank-$d$ truncated Singular Value Decomposition (SVD) of the partial covariance operator. Unlike unconditional covariance estimation, the partial operator involves a residualization term that is not directly observable. The authors formulate this as a bi-level optimization problem:

   • Outer Problem: Learns feature maps $u_\theta(X)$ and $v_\theta(Y, Z)$ to approximate the leading singular functions of the partial operator.
   • Inner Problem: Learns a feature map $w_\theta(Z)$ to implicitly handle the conditioning on $Z$ via a low-rank auxiliary problem derived from the cyclic property of the trace.

   The learning objective is a contrastive loss derived from the Hilbert-Schmidt norm of the operator approximation error, regularized with orthonormality constraints to ensure the learned features are well-conditioned. The algorithm uses neural networks parametrized by $\theta$ and optimizes via stochastic gradient descent on mini-batches.

2. Test Statistic Construction:
   Once the representations $\hat{u}_\theta, \hat{v}_\theta, \hat{w}_\theta$ are learned on a training set, the test statistic is computed on a held-out test set. The statistic is defined as:
   $$\hat{T}_n = n \| \hat{C}_{\hat{U}_\theta \hat{V}_\theta} - \hat{C}_{\hat{U}_\theta \hat{W}_\theta} \hat{C}_{\hat{W}_\theta \hat{V}_\theta} \|_F^2$$
   where $\hat{C}$ denotes empirical cross-covariance matrices of the learned features. This statistic measures the residual dependence between $X$ and $Y$ after accounting for $Z$ in the learned spectral subspace.

3. Theoretical Guarantees:

   • Validity (Null Hypothesis): Under the null hypothesis and mild regularity assumptions (sub-Gaussianity of features), the authors prove that $\hat{T}_n$ converges in distribution to a chi-squared distribution with $d^2$ degrees of freedom, provided the representation learning error ($E_{val}$) vanishes. This allows for the use of standard chi-squared critical values without permutation.
   • Power (Alternative Hypothesis): The power of the test is governed by the quality of the learned representations ($E_{pow}$). The authors establish that the test achieves high power if the signal strength (Hilbert-Schmidt norm of the partial operator) exceeds a threshold determined by the representation error and statistical estimation error.

Key Contributions

1. Algorithmic Innovation: Introduction of a scalable, bi-level contrastive learning algorithm to estimate the leading spectral features of partial covariance operators, overcoming the adaptivity and scalability bottlenecks of classical kernel methods.
2. Theoretical Analysis: A comprehensive theoretical framework linking representation learning error to test performance. The paper provides guarantees for Type I error control (asymptotic chi-squared validity) and power, characterizing the trade-off between signal strength, representation quality, and sample size.
3. Empirical Validation: Extensive experiments on synthetic data (including post-nonlinear models and high-dimensional nonsmooth variants) and real-world data (breast cancer histology and gene expression). The results demonstrate that SpectralCIT maintains robust Type I error control and high power in regimes where kernel-based methods fail or lose validity.

Results

• Synthetic Benchmarks: In post-nonlinear models with varying conditioning dimensions ($d_Z \in [50, 300]$), SpectralCIT consistently controlled Type I error at the nominal level ($\alpha=0.05$) while achieving high power. In contrast, methods like KCIT, RCIT, and GCIT failed to control Type I error in high dimensions, and DGCIT lost control entirely.
• Nonsmooth Data: On a high-dimensional, nonsmooth data model designed to violate the smoothness assumptions of local permutation tests (NNLSCIT), SpectralCIT remained robust, whereas NNLSCIT suffered a complete breakdown in Type I error control.
• Scalability: SpectralCIT was found to be approximately $2\times$ faster than KCIT in high-dimensional settings, as its complexity is dominated by low-rank operations in the truncation dimension $d$ rather than $N \times N$ kernel matrix computations.
• Real-World Application: Applied to breast cancer data (TCGA-BRCA), the method tested whether molecular profiles ($X$) provide prognostic information beyond histological image features ($Z$) for predicting survival ($Y$). While linear models and other CI tests failed to reject the null, SpectralCIT strongly rejected it, detecting complex non-linear dependencies confirmed by an XGBoost model.

Significance and Claims
The paper claims to offer a "principled and statistically grounded path toward scalable CI testing." By bridging kernel-based theory with modern representation learning, SpectralCIT addresses the critical trade-off between validity and scalability. The authors emphasize that their approach does not require restrictive structural assumptions (like Lipschitz continuity) or access to the conditional distribution $P(X|Z)$. Instead, validity is contingent on the quality of the learned spectral representations, a condition that can be controlled via the learning process. The work suggests that learning the spectral structure of dependence operators is a viable strategy for overcoming the "curse of dimensionality" in nonparametric conditional independence testing.