Learning heterogeneous treatment effects under principal stratification

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.

• The paper introduces a novel framework for estimating conditional principal causal effects (CPCEs) under principal stratification using data-adaptive methods.
• The paper develops a sequentially orthogonalized, double cross-fit, doubly robust estimator (PLSS-DCDR) that attains oracle efficiency despite complex nuisance structures.
• The paper demonstrates robust performance through simulations and an ARDS trial application, validating inference under non-monotonicity and varying sensitivity parameters.

Estimating Conditional Heterogeneous Treatment Effects under Principal Stratification

Introduction and Problem Setting

The manuscript "Learning heterogeneous treatment effects under principal stratification" (2606.29076) presents a rigorous framework for data-adaptive estimation and inference of conditional principal causal effects (CPCEs) within the principal stratification framework. Principal stratification is fundamental to causal inference when post-treatment intermediate variables complicate the definition or interpretation of treatment effects, as in noncompliance and truncation by death settings. Traditionally, methodological research focused on average principal causal effects (PCEs) within each latent stratum (e.g., compliers, always-takers). However, the quantification of effect heterogeneity within strata, conditional on effect modifiers, is largely unexplored---despite its high relevance for subgroup or individualized treatment decision-making.

This work advances the field by formalizing a family of CPCEs, which quantify the treatment effect within each principal stratum $G=d_0d_1$ conditional on a vector of effect modifiers $X$. These estimands generalize conventional PCEs and the conditional local average treatment effect (CLATE) but are estimable beyond the settings that assume monotonicity and exclusion restriction. The main technical challenge emerges from the fact that principal strata are unobservable and, without strong assumptions, the interplay of modeling the distribution of latent strata and outcomes is nontrivial.

Identification under Principal Ignorability and Odds Ratio Sensitivity

CPCEs are shown to be identifiable under a set of structural assumptions, most notably principal ignorability (an extension of mean-ignorability to the principal stratification context) and an odds-ratio parameterization linking the unidentifiable latent principal score to observable data. The latter, through a sensitivity function $\theta(C)$, unifies and relaxes (i) monotonicity (as $\theta(C)\to\infty$) and (ii) counterfactual independence ($\theta(C)=1$) as special cases, thus enabling sensitivity analyses relative to these commonly invoked assumptions. The authors derive nonparametric g-computation representations for CPCEs indexed by user-chosen effect modifiers, which is a significant theoretical contribution.

Double Cross-Fit Doubly Robust Machine Learner (DCDR)

Central to the methodology is the proposal of a sequentially orthogonalized, double cross-fit, doubly robust estimator (PLSS-DCDR learner) for the CPCE. The procedure is motivated by the complex nested nuisance structure: base nuisance functions (propensity, conditional mean, principal scores) are required for intermediate pseudo-outcomes, which themselves enter a further pseudo-outcome for the CPCE. The estimation is operationalized using doubly robust orthogonal pseudo-outcomes, with stagewise cross-fitting (outer for base nuisance, inner for intermediate nuisance), and final regression-on-basis via penalized least squares sieves (typically P-splines). This construct ensures:

1. Orthogonality: The CPCE estimation equation is Neyman-orthogonal to nuisance function perturbations, conferring robustness to slow or inconsistent nuisance learning.
2. Double Robustness: Second-order error propagation ensures that the impact of nuisance estimation error is asymptotically negligible under weak regularity.
3. Oracle Efficiency: As $n\to\infty$, the estimator attains the efficiency bound as if the first-stage nuisance functions were known.

The estimator supports a wide class of machine learning approaches in the first and second stages, consistently producing valid inferences as long as nuisance functions converge at fast enough rates.

Figure 1: Simulation results for the root mean integrated squared error (MISE), uniform empirical coverage, and average integrated width of uniform bands for the PLSS-DCDR learner on the complier CPCE under non-monotonicity.

Large-Sample Theory and Inference

The asymptotic theory extends the high-level results for debiased/double machine learning estimators to the case of nonparametric density ratio pseudo-outcomes under principal stratification. Detailed $\mathcal{L}^2$ and $\mathcal{L}^\infty$ rates are established for the PLSS-DCDR estimator, including influence function-based asymptotic linearity, normality at fixed $x$, and strong approximation by a Gaussian process for simultaneous inference.

Critical regularity conditions include sieve basis regularity, boundedness of the odds ratio sensitivity, and undersmoothing. The theory covers both pointwise Wald intervals and honest uniform confidence bands, computed via a fast Gaussian (or Bayesian) bootstrap exploiting the strong Gaussian process approximation.

Figure 2: Estimated CSACE curves for always-survivors with associated pointwise and uniform confidence bands, spanning a range of odds ratio $\theta$ values in the ARDS dataset.

Empirical and Simulation Studies

The simulation results provide strong numerical evidence for the finite-sample properties and robustness of the proposed methodology. When the data-generating process includes principal strata without monotonicity, plug-in or ratio-based DR estimators designed for IV/monotonic settings exhibit severe bias or numerical failure, whereas the PLSS-DCDR estimator achieves nominal coverage and root-MISE across all strata and effect modifiers examined.

Comprehensive sensitivity analyses, both for monotonicity misspecification (i.e., fitting a monotonic model when non-monotonicity holds and vice versa), show highly non-robust behavior of classical estimators and solid performance for the proposed approach when correct assumptions are encoded.

Tables summarizing pointwise coverage, MCSD, and bias (see Table 1 below) demonstrate minimal bias and coverage near nominal levels across all choices of smoothing parameters.

$X$0	$X$1	GCV Bias	GCV CP (%)	GCV.under Bias	Under CP (%)	Unpenalized Bias	Unpenalized CP (%)
5	-0.50	0.02	98.8	0.01	97.7	0.01	96.0
10	0.25	0.02	99.5	-0.01	96.7	-0.01	94.9

Table 1: Selected pointwise inference results for compliers with varying smoothing strategies (non-monotonicity).

Application: ARDS Trial

The method is applied to a high-profile ARDS clinical trial to investigate the conditional survivor average causal effect (CSACE) of high vs. low positive-end expiratory pressure on days-to-return-home, conditioning on driving pressure. Unlike standard IV estimators, which are inapplicable due to clear monotonicity violations (negative principal scores for compliers), the PLSS-DCDR is successfully applied for a range of odds ratios. The results reveal strongly non-monotonic heterogeneity: the benefit of PEEP peaks at moderate driving pressures, is absent at low levels, and reverses to harm in high driving pressures, consistent across all sensitivity scenarios.

Figure 2: Estimated CSACE curves and confidence bands for various sensitivity parameterizations, demonstrating effect heterogeneity and monotonic deviation in the ARDS trial.

Simulation-Based Comparison: CPCE across all Principal Strata

Rich simulation studies investigate the MISE, coverage, and empirical band width for CPCEs in all principal strata (always-takers, never-takers, defiers, compliers). Results are consistent: the PLSS-DCDR estimator is robustly well-calibrated unless strong assumptions (e.g., monotonicity) are violated in the data-generating process and incorrectly imposed in estimation, which produces pathologically high bias and invalid inference for ratio-based estimators.

Figure 3: Root mean integrated squared error (MISE), coverage probability, and uniform band width for always-takers, illustrating stable performance of the PLSS-DCDR estimator under non-monotonicity.

Figure 4: Analogous metrics (MISE, coverage, band width) for never-takers, underscoring the robustness of the estimator.

Implications and Future Directions

The technical implications are multifold:

• The framework extends the scope of principal stratification analysis to cover arbitrary principal strata and heterogeneous effect modification, without relying on exclusion or monotonicity as point-identifying assumptions.
• The odds-ratio sensitivity parameter enables transparent and practical sensitivity analysis with a parameter familiar from epidemiology.
• The proposed DCDR estimator resolves estimation challenges for nested and nonlinear nuisance functions by leveraging sequential orthogonalization and double cross-fitting, achieving optimal bias-variance tradeoffs even with generic ML function classes.

Practically, the method enables valid, interpretable estimates of treatment effect heterogeneity in complex designs (e.g., post-randomization selection, noncompliance, truncation by death) where naive or classical IV-based approaches would either be inapplicable or badly miscalibrated.

Further research directions include full sensitivity analyses for departures from principal ignorability, extensions to multiple intermediate outcomes, and integration of more flexible machine learning base learners into the DCDR framework. The theoretical tools developed here (oracle-efficient learning with nested orthogonalization) may also have broad implications for other domains in semi- and nonparametric causal inference.

Conclusion

This paper establishes a unified framework for efficient, robust estimation and inference of heterogeneous CPCEs under principal stratification, equipped with valid uncertainty quantification for both pointwise and uniform inference. The key innovations---odds-ratio sensitivity identification, sequential orthogonal debiasing, and doubly robust, double cross-fit learning---yield estimators that remain valid under a wide range of identification assumptions, thus filling a critical gap in both methodological and applied causal inference.

Reference:

"Learning heterogeneous treatment effects under principal stratification" (2606.29076)