Causal Stability Selection

Abstract: Identifying covariates that modify treatment effects is a central problem in causal inference. Yet existing data-adaptive procedures do not provide finite-sample control over the expected number of false discoveries, risking spurious findings that fail to replicate. We introduce causal stability selection, an algorithm that combines cross-fitted estimation of conditional average treatment effects with integrated path stability selection. The method accommodates arbitrary treatment effect estimators and arbitrary base selectors, and produces a selection set with an explicit, non-asymptotic bound on the expected number of false positives. Under standard causal identifying assumptions and regularity conditions on the base selector, we prove that the estimated selection probabilities converge to their oracle counterparts at the rate of the underlying treatment effect estimator. This establishes a direct connection between treatment effect estimation and effect modifier discovery. We illustrate the method on a randomized trial in oncology and on observational data on maternal smoking and infant birthweight.

• The paper introduces CausalStabSel, a novel method that integrates cross-fitted CATE estimation with stability selection to discover treatment effect modifiers.
• It provides non-asymptotic guarantees by controlling the expected number of false positives while maintaining strong detection power in diverse data scenarios.
• Empirical evaluations on simulations and real-world data confirm its robustness under confounding, nonlinearity, and high-dimensional conditions.

Causal Stability Selection for Effect Modifier Discovery

Introduction and Motivation

Effect modifier discovery is fundamental in causal inference: it aims to identify covariates that modulate a treatment’s effect. While significant advances have occurred in data-adaptive procedures for discovering treatment effect heterogeneity, these methods typically lack finite-sample guarantees on the expected number of false discoveries. This gap is consequential, as inflated false positive rates undermine replicability and the reliability of precision interventions.

The paper "Causal Stability Selection" (2605.09300) proposes a method—Causal Stability Selection (CausalStabSel)—that enables effect modifier selection with explicit, non-asymptotic upper bounds on false positives. The approach integrates cross-fitted estimation of conditional average treatment effects (CATE) with integrated path stability selection, thereby achieving error quantifiability and modularity with respect to both CATE estimation and selection mechanisms.

Problem Definition and Framework

The effect modifier discovery task is formalized under the potential outcomes framework. For observed data $(x_i, y_i, z_i)$, the CATE $\tau(x)$ and the baseline outcome function $\mu(x)$ are defined respectively as:

$$\tau(x) = \mathbb{E}[Y(1) - Y(0) \mid X = x], \quad \mu(x) = \mathbb{E}[Y(0) \mid X = x].$$

A variable $X_j$ is an effect modifier if $\tau(X) \not\!\perp\!\!\!\perp X_j \mid X_{-j}$, i.e., $X_j$ is associated with residual heterogeneity in $\tau$ given the other covariates. The goal is to estimate the set $\mathcal{E}$ of effect modifiers with (non-asymptotic) control on the expected number of false selections.

Methodology: CausalStabSel Algorithm

CausalStabSel generalizes stability selection to the causal effect heterogeneity context. Traditional stability selection aggregates results from multiple resamples but assumes direct observability of the outcome of interest. In the effect modifier setting, the unobserved $\tau(x_i)$ must be estimated, introducing additional uncertainty and bias.

CausalStabSel addresses this by:

• Employing repeated cross-fitted estimation: for each random split, one subsample is held out for CATE estimation, then variables are selected on the remaining samples using these out-of-sample $\tau(x)$0 predictions.
• Estimating selection probabilities $\tau(x)$1 across splits and penalization levels $\tau(x)$2.
• Final selection is achieved via integrated path stability selection (IPSS), aggregating over $\tau(x)$3 and utilizing a cubic transformation to improve sharpness and power relative to traditional maximum-based selection criteria.

This design separates estimation of $\tau(x)$4 from variable selection, suppressing overfitting and spurious discovery. The approach accommodates arbitrary CATE estimators and selection algorithms, including penalized regression and importance-score thresholding.

Theoretical Guarantees

The paper establishes several key results:

• Unbiasedness: Estimated selection probabilities converge to their oracle (true-$\tau(x)$5) counterparts at the same rate as the underlying CATE estimator. Under mild regularity (Lipschitz importance scores, bounded densities), the convergence is $\tau(x)$6 with parametric-rate CATE estimators.
• Oracle Consistency: The selected set from CausalStabSel converges in expected symmetric difference to the oracle selection set, provided the cross-fitting subsample size grows sublinearly with total sample size.
• Error Control: A tight, non-asymptotic upper bound for the expected number of false positives is derived, and thresholding by expected false positives (efp-score) achieves target FDR levels by construction.

Empirical Evaluation

CausalStabSel is validated through comprehensive simulations and real-data applications. The simulation suite encompasses varying confounding, sample size, covariate dimension, correlation, signal-to-noise ratio, and prognostic strength—both in linear and nonlinear generative regimes.

In an illustrative example, CausalStabSel is contrasted with alternative effect modifier selection strategies (Integrated Path Stability Selection, Knockoffs, Benjamini-Hochberg, and Lasso) using cross-fitted or split-sample pseudo-outcomes. Only CausalStabSel consistently controls the false discovery rate at or below nominal levels across a spectrum of conditions, without sacrificing detection power.

Figure 1: Comparative TPR and FDR results for effect modifier discovery methods, highlighting that only CausalStabSel achieves nominal FDR control under practical data conditions.

Further experiments dissect the robustness of FDR control and sensitivity to confounding and nonlinearities:

Figure 2: CausalStabSel provides valid FDR control and strong TPR across increasing levels of confounding in linear settings.

Figure 3: Consistent FDR control by CausalStabSel in nonlinear settings, demonstrating broad robustness.

These empirical findings are consistent under variations in the size of the covariate set, number of effect modifiers, and other key design parameters.

Real-World Applications

The practical efficacy of CausalStabSel is assessed on two datasets:

• PRIME Oncology RCT: CausalStabSel successfully recovers known clinical effect modifiers (e.g., KRAS mutation status) and plausible clinical predictors, outperforming simple thresholding and classical regression in terms of replicability and parsimony.
• Maternal Smoking / Birthweight Observational Study: CausalStabSel identifies established modifiers (e.g., maternal age, parity) with supporting evidence from held-out inference.

These results underscore the utility of CausalStabSel in both randomized and observational designs, with implications for subgroup targeting and personalized intervention strategies.

Implications and Future Directions

CausalStabSel delivers a practical variable selection algorithm with proven false positive control for discovery of treatment effect modifiers, making it suitable for high-dimensional scenarios and heterogeneous data regimes. The explicit link between effect modifier estimation and CATE accuracy reinforces the importance of advanced CATE estimators in high-stakes scientific or policy settings.

Potential extensions include:

• Adapting CausalStabSel for subgroup discovery and time-varying treatments.
• Employing the framework for causal inference in complex structures, including interference and networked data.
• Developing post-selection inference techniques that leverage CausalStabSel’s design to provide valid uncertainty quantification.

Cross-fitted stability selection is also conceptually aligned with best practices in modern machine learning for regularization and out-of-sample validation, suggesting broader applicability beyond causal effect heterogeneity.

Conclusion

CausalStabSel formalizes and addresses the challenge of effect modifier discovery with finite-sample guarantees and modularity in estimator and selector specification. Extensive empirical and theoretical results confirm that it uniquely achieves both valid error control and strong detection power across linear, nonlinear, confounded, and high-dimensional settings. Its design is directly extensible to other areas of causal inference requiring robust, reproducible, and interpretable effect heterogeneity analysis.