Learning heterogeneous treatment effects under principal stratification

2026-06-27

AI Summary

This study addresses the limitation of existing principal stratification methods, which primarily focus on average causal effects and struggle to capture heterogeneity in individualized treatment effects within strata. Moving beyond the monotonicity assumption, the authors identify and estimate conditional principal causal effects by integrating principal ignorability with a sensitivity analysis parameterized via odds ratios. To this end, they propose a novel doubly cross-fitted, doubly robust machine learning estimator that combines sequential orthogonal learning with regularized least squares sieve methods, enabling efficient estimation under nested nuisance structures. Theoretical analysis establishes $L^2$ convergence and asymptotic normality, demonstrating oracle efficiency and facilitating the construction of uniform confidence bands. Simulations and an application to a clinical trial on acute lung injury corroborate the method’s superior performance.

Abstract

Principal stratification provides a foundational framework for causal inference with intermediate outcomes by defining causal effects within subpopulations, yet existing work has largely focused on average effects across strata rather than treatment effect heterogeneity within strata. Such within-stratum heterogeneity informs individualized treatment decisions but the associated methods are sparse. We address this gap by studying the identification and estimation of the conditional principal causal effects under principal ignorability combined with an odds ratio sensitivity parameterization, which relaxes the monotonicity assumption. To efficiently learn these estimands, we propose a novel doubly cross-fit doubly robust machine learner that resolves the nested nuisance structure inherent to principal stratification. Leveraging sequential orthogonal learning with regularized least-squares sieves, we derive $\mathcal{L}^2$ and uniform limit theory, establish oracle efficiency, and construct uniform confidence bands for the proposed estimator. We use simulations to demonstrate the finite-sample performance of our estimator, and provide an empirical analysis of a randomized trial in acute lung injury, revealing informative patterns of treatment effect heterogeneity within the always-survivor subpopulation.

Problem

• heterogeneous treatment effects
• principal stratification
• causal inference
• within-stratum heterogeneity
• individualized treatment decisions

Innovation

• principal stratification
• treatment effect heterogeneity
• doubly robust estimation
• cross-fitting
• sieve estimation