Dimension-reduced outcome-weighted learning for estimating individualized treatment regimes in observational studies

TL;DR

This work introduces a gradient kernel dimension reduction (gKDR) framework to identify a low-dimensional central mean subspace that preserves treatment-effect heterogeneity in observational data. By integrating kernel-based covariate balancing (KCB) with SDR, the method allows treatment to depend on full covariates while targeting a reduced representation for learning optimal ITRs through augmented outcome-weighted learning (AOL). Theoretical guarantees include universal consistency of the resulting decision rule, convergence of balancing weights to inverse propensity weights, and consistency of the estimated subspace and risk. Empirically, the approach (DOL) demonstrates superior finite-sample performance to competing methods in simulations and yields improved modified value functions in a real ICU sepsis dataset, suggesting meaningful gains in clinical decision-making with observational data. The framework is scalable, flexible to high-dimensional covariates, and adaptable to observational settings, with potential extensions to missing data and multi-valued or dynamic treatments.

Abstract

Individualized treatment regimes (ITRs) aim to improve clinical outcomes by assigning treatment based on patient-specific characteristics. However, existing methods often struggle with high-dimensional covariates, limiting accuracy, interpretability, and real-world applicability. We propose a novel sufficient dimension reduction approach that directly targets the contrast between potential outcomes and identifies a low-dimensional subspace of the covariates capturing treatment effect heterogeneity. This reduced representation enables more accurate estimation of optimal ITRs through outcome-weighted learning. To accommodate observational data, our method incorporates kernel-based covariate balancing, allowing treatment assignment to depend on the full covariate set and avoiding the restrictive assumption that the subspace sufficient for modeling heterogeneous treatment effects is also sufficient for confounding adjustment. We show that the proposed method achieves universal consistency, i.e., its risk converges to the Bayes risk, under mild regularity conditions. We demonstrate its finite sample performance through simulations and an analysis of intensive care unit sepsis patient data to determine who should receive transthoracic echocardiography.

Figures (9)

• Figure 1: Blue points correspond to the treatment group "+1" and red points correspond to the treatment group "-1". The sample sizes for the +1 and -1 treatment groups are 110 and 90, respectively. (a) shows the distribution of the data with respect to the covariates $\left(X^{(1)}, X^{(2)}\right)$. The magnitudes of the outcomes are represented by varying point sizes. The solid line represents the true projection line that best captures treatment effect heterogeneity, while the dashed line denotes the estimated projection line obtained from our proposed SDR approach. (b) displays the true pseudo-outcomes ($Y/\mathrm{Pr}(+1|X)$ and $-Y/\mathrm{Pr}(-1|X)$ for treatment groups $A=+1$ and $A=-1$) in the reduced subspace, where $\left(X^{(1)}, X^{(2)}\right)$ is projected onto $V = B_0^\top X$. The curve represents the heterogeneous treatment effect as a function of $V$. (c) shows the pseudo-outcomes projected onto $V^{\perp}$, the space perpendicular to that spanned by $B_0$. The solid line shows the treatment effect as a function of $V^\perp$. (d) shows the pseudo-outcomes as a function of $\hat{V} = \hat{B}^\top X$, where $\hat{B}$ is an estimate of $B_0$ obtained using the proposed SDR framework. The solid line shows the true heterogeneous treatment effect along the different levels of the dimension-reduced covariate. The dashed line shows the estimated decision rule along the different levels of the dimension-reduced covariate. The signs of these functions, which determine the true and estimated treatment assignments, are equal in the range $\hat{V} \in (-1, 1)$, where approximately 95% of the data points lie.
• Figure 2: Scatterplot of the Bayes optimal treatment regimes for different values of the covariates generated under Setting 1. The red circles and blue crosses represent data points from the test dataset, corresponding to $d^*(X) = -1$ (or $d^*(V_0) = -1$) and $d^*(X) = +1$ (or $d^*(V_0) = +1$), respectively. The solid lines in the left plot display the decision boundaries in the dimension-reduced space, defined as $\left\{v = \left(v^{(1)}, v^{(2)}\right) : \tilde{f}^*(v) = 0 \right\}$.
• Figure 3: Simulation results for $n = 1{,}000$ from the randomized scenarios across the three different settings. Accuracy represents the proportion of correctly predicted optimal treatments. Value (%) represents the value function recovered by the estimated decision rule as a percentage of the Bayes optimal value function. The performance of the proposed method is shown under DOL-O and DOL-L: DOL-O uses the oracle $g(x)$, while DOL-L uses $g_{\hat{w}}(x)$ estimated via linear regression.
• Figure 4: Simulation results for $n = 1{,}000$ from the non-randomized scenarios across the three different settings. Accuracy represents the proportion of correctly predicted optimal treatments. Value (%) represents the value function recovered by the estimated decision rule as a percentage of the Bayes optimal value function. The performance of the proposed method is shown under KCB-O and KCB-L: KCB-O uses the oracle $g(x)$, while KCB-L uses the fitted values of $g_{\hat{w}}(x)$ obtained via linear regression. IPW-O uses the oracle $g(x)$, and IPW-L uses $g_{\hat{w}}(x)$ from linear regression, with propensity scores estimated via logistic regression.
• Figure 5: Modified value functions over the 100 repeats. DOL represents the proposed method. AOL used the raw patient characteristics for learning ITR and fitting logistic regression to estimate IPW. QL is the $\ell_1 \text{-} PLS$. "Treat all" represents the average effect of TTE. P-values from paired $t$-tests are displayed using asterisks: **** $<0.0001$; *** $<0.001$; * $< 0.05$.

Theorems & Definitions (14)

• Remark 1
• Theorem 1
• Corollary 1
• Theorem 2
• Theorem 3
• Theorem 4
• Corollary 2
• Lemma 1
• Proposition 1
• Theorem 1