Recovering Sparse and Interpretable Subgroups with Heterogeneous Treatment Effects with Censored Time-to-Event Outcomes

Data driven decision making across multiple disciplines including healthcare, epidemiology, econometrics and prognostics often involves establishing efficacy of an intervention when outcomes are measured in terms of the time to an adverse event, such as death, failure or onset of a critical condition. Typically the analysis of such studies involves assigning a patient population to two or more different treatment arms often called the ‘treated’ (or ‘exposed’) group and the ‘control’ (or ‘placebo’) group and observing whether the populations experience an adverse event (for instance death or onset of a disease) over the study period at a rate that is higher (or lower) than for the control group. Efficacy of a treatment is thus established by comparing the relative difference in the rate of event incidence between the two arms called the hazard ratio. However, not all individuals benefit equally from an intervention. Thus, very often potentially beneficial interventions are discarded even though there might exist individuals who benefit, as the population level estimates of treatment efficacy are inconclusive.

In this paper we assume that patient responses to an intervention are typically heterogeneous and there exists patient subgroups that are unaffected by (or worse, harmed) by the intervention being assessed. The ability to discover or phenotype these patients is thus clinically useful as it would allow for more precise clinical decision making by identifying individuals that actually benefit from the intervention being assessed.

Towards this end, we propose Sparse Cox Subgrouping, (SCS) a latent variable approach to model patient subgroups that demonstrate heterogeneous effects to an intervention. As opposed to existing literature in modelling heterogeneous treatment effects with censored time-to-event outcomes our approach involves structured regularization of the covariates that assign individuals to subgroups leading to parsimonious models resulting in phenogroups that are interpretable. We release a python implementation of the proposed SCS approach as part of the auton-survival package (Nagpal et al., 2022b) for survival analysis:

https://autonlab.github.io/auton-survival/

Large studies especially in clinical medicine and epidemiology involve outcomes that are time-to-events such as death, or an adverse clinical condition like stroke or cancer. Treatment efficacy is typically estimated by comparing event rates between the treated and control arms using the Proportional Hazards (Cox, 1972) model and its extensions.

Identification of subgroups in such scenarios has been the subject of a large amount of traditional statistical literature. Large number of such approaches involve estimation of the factual and counterfactual outcomes using separate regression models (T-learner) followed by regressing the difference between these estimated potential outcomes. Within this category of approaches, Lipkovich et al. (2011) propose the subgroup identification based on differential effect search (SIDES) algorithm, Su et al. (2009) propose a recursive partitioning method for subgroup discovery, Dusseldorp and Mechelen (2014) propose the qualitative interaction trees (QUINT) algorithm, and Foster et al. (2011) propose the virtual twins (VT) method for subgroup discovery involving decision tree ensembles. We include a parametric version of such an approach as a competing baseline.

Identification of heterogeneous treatment effects (HTE) is also of growing interest to the machine learning community with multiple approaches involving deep neural networks with balanced representations (Shalit et al., 2017; Johansson et al., 2020), generative models Louizos et al. (2017) as well as Non-Parametric methods involving random-forests (Wager and Athey, 2018) and Gaussian Processes (Alaa and Van Der Schaar, 2017). There is a growing interest in estimating HTEs from an interpretable and trustworthy standpoint (Lee et al., 2020; Nagpal et al., 2020; Morucci et al., 2020; Wu et al., 2022; Crabbé et al., 2022). Wang and Rudin (2022) propose a sampling based approach to discovering interpretable rule sets demonstrating HTEs.

However large part of this work has focused extensively on outcomes that are binary or continuous. The estimation of HTEs in the presence of censored time-to-events has been limited. Xu et al. (2022) explore the problem and describe standard approaches to estimate treatment effect heterogeneity with survival outcomes. They also describe challenges associated with existing risk models when assessing treatment effect heterogeneity in the case of cardiovascular health.

There has been some initial attempt to use neural network for causal inference with censored time-to-event outcomes. Curth et al. (2021) propose a discrete time method along with regularization to match the treated and control representations. Chapfuwa et al. (2021)’s approach is related and involves the use of normalizing flows to estimate the potential time-to-event distributions under treatment and control. While our contributions are similar to Nagpal et al. (2022a), in that we assume treatment effect heterogeneity through a latent variable model, our contribution differs in that 1) Our approach is free of the expensive Monte-Carlo sampling procedure and 2) Our generalized EM inference procedure allows us to naturally incorporate structured sparsity regularization, which helps recovers phenogroups that are parsimonious in the features they recover that define subgroups.

Survival and time-to-event outcomes occur pre-eminently in areas of cardiovascular health. One such area is reducing combined risk of adverse outcomes from atherosclerotic disease¹¹1A class of related clinical conditions from increasing deposits of plaque in the arteries, leading to Stroke, Myorcardial Infarction and other Coronary Heart Diseases. (Herrington et al., 2016; Furberg et al., 2002; Group, 2009; Buse et al., 2007) The ability of recovering groups with differential benefits to interventions can thus lead to improved patient care through framing of optimal clinical guidelines.

Notation

As is standard in survival analysis, we assume that we either observe the true time-to-event or the time of censoring $U=\mathrm{min}\{T,C\}$ indicated by the censoring indicator defined as $\Delta=\bm{1}\{T<C\}$. We thus work with a dataset of right censored observations in the form of 4-tuples, $\mathcal{D}=\{({\bm{x}}_{i},\delta_{i},{\bm{u}}_{i},{\bm{a}}_{i})\}_{i=1}^{n}$, where ${\bm{u}}_{i}\in\mathbb{R}^{+}$ is the time-to-event or censoring as indicated by $\delta_{i}\in\{0,1\}$, ${\bm{a}}_{i}\in\{0,1\}$ is the indicator of treatment assignment, and ${\bm{x}}_{i}$ are individual covariates that confound the treatment assignment and the outcome.

Assumption 1 (Independent Censoring)

The time-to-event $T$ and the censoring distribution $C$ are independent conditional on the covariates $X$ and the intervention $A$.

Model

Consider a maximum likelihood approach to model the data $\mathcal{D}$ the set of parameters $\bm{\Omega}$. Under Assumption 1 the likelihood of the data $\mathcal{D}$ can be given as,

$\displaystyle\mathcal{L}(\bm{\Omega};\mathcal{D})$

$\displaystyle\propto\prod_{i=1}^{|\mathcal{D}|}\bm{\lambda}(u_{i}|X={\bm{x}}_{i},A={\bm{a}}_{i})^{\delta_{i}}\bm{S}(u_{i}|X={\bm{x}}_{i},A={\bm{a}}_{i}),$

(1)

here $\bm{\lambda}(t)=\mathop{\lim}\limits_{\Delta t\to 0}\frac{\mathbb{P}(t\leq T<t+\Delta t|T\geq t)}{\Delta t}$ is the hazard and $\bm{S}(t)=\mathbb{P}(T>t)$ is the survival rate.

Assumption 2 (PH)

The distribution of the time-to-event $T$ conditional on the covariates and the treatment assignment obeys proportional hazards.

From Assumption 2 (Proportional Hazards), an individual with covariates $(X={\bm{x}})$ under intervention $(A={\bm{a}})$ under a Cox model with parameters $\beta$ and treatment effect $\omega$ is given as

$\displaystyle\bm{\lambda}({\bm{t}}|A={\bm{a}},X={\bm{x}})=\bm{\lambda}_{0}(t)\exp\big{(}\bm{\beta}^{\top}{\bm{x}}+\bm{\omega}\cdot\bm{a}\big{)},$

(2)

Here, $\bm{\lambda}_{0}(\cdot)$ is an infinite dimensional parameter known as the base survival rate. In practice in the Cox’s model the base survival rate is a nuisance parameter and is estimated non-parametrically. In order to model the heterogeneity of treatment response. We will now introduce a latent variable $Z\in\{0,1,-1\}$ that mediates treatment response to the model,

	$\displaystyle\bm{\lambda}({\bm{t}}|A={\bm{a}},X={\bm{x}},Z={\bm{k}})$	$\displaystyle=\bm{\lambda}_{0}(t)\exp(\beta^{\top}{\bm{x}})\exp(\bm{\omega})^{\bm{ka}},$
	$\displaystyle\text{and, }\;\;\mathbb{P}(Z$	$\displaystyle={\bm{k}}|X={\bm{x}})=\frac{\exp(\bm{\theta}_{k}^{\top}{\bm{x}})}{\sum_{j}\exp(\bm{\theta}_{j}^{\top}{\bm{x}})}.$		(3)

Here, $\bm{\omega}\in\mathbb{R}$ is the treatment effect, and $\bm{\theta}\in\mathbb{R}^{k\times d}$ is the set of parameters that mediate assignment to the latent group $Z$ conditioned on the confounding features ${\bm{x}}$. Note that the above choice of parameterization naturally enforces the requirements from the model as in Figure 1. Consider the following scenarios,

Case 1: The study population consists of two sub-strata ie. $Z\in\{0,+1\}$, that are benefit and are unaffected by treatment respectively.

Case 2: The study population consists of three sub-strata ie. $Z\in\{0,+1,-1\}$, that benefit, are harmed or unaffected by treatment respectively.

Following from Equations 1 & 3, the complete likelihood of the data $\mathcal{D}$ under this model is,

	$\displaystyle\mathcal{L}(\bm{\Omega};\mathcal{D})=\prod_{i=1}^{|\mathcal{D}|}\sum_{k\in Z}\bigg{(}\bm{\lambda}_{0}(u_{i})\bm{h}({\bm{x}},{\bm{a}},{\bm{k}})\bigg{)}^{\delta_{i}}$	$\displaystyle{\bm{S}}_{0}(u_{i})^{\bm{h}({\bm{x}},{\bm{a}},{\bm{k}})}\mathbb{P}(Z=k|X={\bm{x}}_{i})$
	$\displaystyle\text{ where, }\ln\bm{h}({\bm{x}},{\bm{a}},{\bm{k}})=\bm{\beta}^{\top}{\bm{x}}+{\bm{k}}\cdot\bm{a}\cdot\bm{w}\text{ and }\ln$	$\displaystyle{\bm{S}}_{0}(\cdot)=-\bm{\Lambda}_{0}(\cdot),$		(4)

Note that $\bm{\Lambda}_{0}(\cdot)=\int_{0}^{t}\bm{\lambda}_{0}(\cdot)$ is the infinite dimensional cumulative hazard and is inferred when learning the model. We will notate the set of all learnable parameters as $\bm{\Omega}=\{\bm{\theta},\bm{\beta},\bm{w},\bm{\Lambda}_{0}\}$.

Shrinkage

In retrospective analysis to recover treatment effect heterogeneity a natural requirement is parsimony of the recovered subgroups in terms of the covariates to promote model interpretability. Such parsimony can be naturally enforced through appropriate shrinkage on the coefficients that promote sparsity. We want to recover phenogroups that are ‘sparse’ in $\bm{\theta}$. We enforce sparsity in the parameters of the latent $Z$ gating function via a group $\ell_{1}$ (Lasso) penalty. The final loss function to be optimized including the group sparsity regularization term is,

	$\displaystyle\mathcal{L}(\bm{\Omega};\mathcal{D})+$	$\displaystyle\bm{\epsilon}\cdot\mathcal{R}(\bm{\theta})\text{ where, }\mathcal{R}(\bm{\theta})=\sum_{d}\sqrt{\sum_{k\in\mathcal{Z}}\big{(}\bm{\theta}^{k}_{d}\big{)}^{2}}$
	and	$\displaystyle\bm{{\epsilon}}>0\text{ is the strength of the shrinkage parameter}.$		(5)

Identifiability

Further, to ensure identifiability we restrict the gating parameters for the $(Z=0)$ to be $0$. Thus $\bm{\theta}_{1}=0$.

Inference

We will present a variation of the Expectation Maximization algorithm to infer the parameters in Equation 3. Our approach differs from Nagpal et al. (2022a, 2021) in that it does not require storchastic Monte-Carlo sampling. Further, our generalized EM inference allows for incorporation of the structured sparsity in the M-Step.

A Semi-Parametric $Q(\cdot)$

Note that the likelihood in Equation 3 is semi-parametric and consists of parametric components and the infinite dimensional base hazard $\bm{\Lambda}(\cdot)$. We define the $Q(\cdot)$ as:

$\displaystyle Q(\bm{\Omega};\mathcal{D})$

$\displaystyle=\sum_{i=1}^{n}\sum_{k\in\mathcal{Z}}\bm{\gamma}^{k}_{i}\bigg{(}\ln\bm{p}_{\bm{\theta}}(Z=k|X={\bm{x}}_{i})+\ln\bm{p}_{\bm{w},\bm{\beta},\bm{\Lambda}}(T|Z=k,X={\bm{x}}_{i})\bigg{)}+\mathcal{R}(\bm{\theta})$

The E-Step

Requires computation of the posteriors counts $\bm{\gamma}:=\bm{p}(Z={k}|T,X=\bm{x},A={\bm{a}}).$

Result 1 (Posterior Counts)

The posterior counts $\bm{\gamma}$ for the latent $Z$ are estimated as,

	$\displaystyle\bm{\gamma}^{k}$	$\displaystyle=\widehat{\mathbb{P}}(Z=k|X={\bm{x}},A={\bm{a}},{\bm{u}})$
		$\displaystyle=\frac{\mathbb{P}({\bm{u}}|Z={\bm{k}},X={\bm{x}},A={\bm{a}})\mathbb{P}(Z={\bm{k}}|X={\bm{x}})}{\sum_{k}\mathbb{P}({\bm{u}}|Z={\bm{k}},X={\bm{x}},A={\bm{a}})\mathbb{P}(Z={\bm{k}}|X={\bm{x}})}$
		$\displaystyle=\frac{{{\bm{h}}({\bm{x}},{\bm{a}},{\bm{k}})}^{\delta_{i}}\widehat{{\bm{S}}}_{0}({\bm{u}})^{{\bm{h}}({\bm{x}},{\bm{a}},{\bm{k}})}\exp(\bm{\theta_{k}}^{\top}\bm{x})}{\mathop{\sum}_{j\in\mathcal{Z}}{{\bm{h}}({\bm{x}},{\bm{a}},{\bm{j}})}^{\delta_{i}}\widehat{{\bm{S}}}_{0}({\bm{u}})^{{\bm{h}}({\bm{x}},{\bm{a}},{\bm{j}})}\exp(\bm{\theta_{j}}^{\top}\bm{x})}.$		(6)

For a full discussion on derivation of the $Q(\cdot)$ and the posterior counts please refer to Appendix B

The M-Step

Involves maximizing the $Q(\cdot)$ function. Rewriting the $Q(\cdot)$ as a sum of two terms,

$\displaystyle Q(\bm{\Omega})=\underbrace{\sum_{i=1}^{n}\sum_{k\in\mathcal{Z}}\bm{\gamma}^{k}_{i}\ln\bm{p}_{\bm{w},\bm{\beta},\bm{\Lambda}_{0}}(T|Z=k,X={\bm{x}}_{i},A={\bm{a}}_{i})}_{{\bm{A}}(\bm{w},\bm{\beta},\bm{\Lambda}_{0})}+\underbrace{\sum_{i=1}^{n}\sum_{k\in\mathcal{Z}}\bm{\gamma}^{k}_{i}\ln\bm{p}_{\bm{\theta}}(Z=k|X={\bm{x}}_{i})+\mathcal{R}(\bm{\theta})}_{{\bm{B}}(\bm{\theta})}$

Result 2 (Weighted Cox model)

The term ${\bm{A}}$ can be rewritten as a weighted Cox model and thus optimized using the corresponding ‘partial likelihood’,

Updates for $\{\bm{\beta},\bm{\omega}\}$: The partial-likelihood, $\mathcal{PL(\cdot)}$ under sampling weights (Binder, 1992) is

$\displaystyle\mathcal{PL}(\bm{\Omega};\mathcal{D})=\sum_{\begin{subarray}{c}i=1,\delta_{i}=1\end{subarray}}^{n}\sum_{k\in\mathcal{Z}}\bm{\gamma}^{k}_{i}\bigg{(}\ln\bm{h}_{k}({\bm{x}}_{i},\bm{a}_{i},{\bm{k}})-\ln$

$\displaystyle\sum_{j\in{\mathsf{RiskSet}(u_{i})}}\sum_{k\in\mathcal{Z}}\bm{\gamma}^{k}_{j}\bm{h}_{k}({\bm{x}}_{j},\bm{a}_{j},{\bm{k}})\bigg{)}\bigg{]}$

(7)

Here $\mathsf{RiskSet}(\cdot)$ is the ‘risk set’ or the set of all individuals who haven’t experienced the event till the corresponding time, i.e. $\mathsf{RiskSet}(t):=\{i:u_{i}>t\}$. $\mathcal{PL(\cdot)}$ is convex in $\{\bm{\beta},\bm{\omega}\}$ and we update these with a gradient step.

Updates for $\bm{\Lambda}_{0}$: The base hazard $\bm{\Lambda}_{0}$ are updated using a weighted Breslow’s estimate (Breslow, 1972; Lin, 2007) assuming the posterior counts $\bm{\gamma}$ to be sampling weights (Chen, 2009),

$\displaystyle\widehat{\bm{\Lambda}}_{0}(t)^{+}=\sum_{i=1}^{n}\sum_{k\in\mathcal{Z}}\bm{1}\{u_{i}<t\}\frac{\bm{\gamma}^{k}_{i}\cdot\delta_{i}}{\mathop{\sum}\limits_{j\in\mathsf{RiskSet}(u_{i})}\mathop{\sum}\limits_{k\in\mathcal{Z}}\bm{\gamma}^{k}_{j}\bm{h}_{k}({\bm{x}}_{j},\bm{a}_{j},{\bm{k}})}$

(8)

Term ${\bm{B}}$ is a function of the gating parameters $\bm{\theta}$ that determine the latent assignment $Z$ along with sparsity regularization. We update ${\bm{B}}$ using a Proximal Gradient update as is the case with Iterative Soft Thresholding (ISTA) for group sparse $\ell_{1}$ regression.

Updates for $\bm{\theta}$: The proximal update for $\bm{\theta}$ including the group regularization (Friedman et al., 2010) term $\mathcal{R}(\cdot)$ is,

$\displaystyle\widehat{\bm{\theta}}^{+}=\mathsf{prox}_{\eta{\epsilon}}\bigg{(}\bm{\theta}-\frac{d}{d\bm{\theta}}{\bm{B}}(\bm{\theta})\bigg{)},\;\;\;\text{ where }\mathsf{prox}_{\eta{\epsilon}}({\bm{y}})=\frac{{\bm{y}}}{||{\bm{y}}||_{2}}\mathrm{max}\{0,||\bm{y}||_{2}-\eta\bm{\epsilon}\}.$

(9)

[!h] Parameter Learning for SCS with a Generalized EM \SetAlgoLined\SetKwInOutInputInput \SetKwInOutreturnReturn \InputTraining set, $\mathcal{D}=\{({\bm{x}}_{i},u_{i},a_{i},\delta_{i})_{i=1}^{n}\}$; maximum EM iterations, $B$, step size $\eta$  

<not converged> \For$b\in\{1,2,...,B\}$ E-Step .

$\bm{\gamma}_{i}^{k}=\frac{{{\bm{h}}({\bm{x}},{\bm{a}},{\bm{k}})}^{\delta_{i}}\widehat{{\bm{S}}}_{0}({\bm{u}})^{{\bm{h}}({\bm{x}},{\bm{a}},{\bm{k}})}\exp(\bm{\theta_{k}}^{\top}\bm{x})}{\mathop{\sum}_{j\in\mathcal{Z}}{{\bm{h}}({\bm{x}},{\bm{a}},{\bm{j}})}^{\delta_{i}}\widehat{{\bm{S}}}_{0}({\bm{u}})^{{\bm{h}}({\bm{x}},{\bm{a}},{\bm{j}})}\exp(\bm{\theta_{j}}^{\top}\bm{x})}$ $\hfill\triangleright$ Compute posterior counts (Equation 6).

M-Step .

$\widehat{\bm{\beta}}^{+}\leftarrow\widehat{\bm{\beta}}-\eta\cdot\nabla_{\bm{\beta}}\mathcal{PL}(\bm{\beta},{\bm{w}};\mathcal{D})$

$\widehat{\bm{w}}^{+}\leftarrow\widehat{\bm{w}}-\eta\cdot\nabla_{\bm{w}}\mathcal{PL}(\bm{\beta},{\bm{w}};\mathcal{D})$ $\triangleright$ Gradient descent update.

$\widehat{\bm{\Lambda}}_{0}(t)^{+}\leftarrow\mathop{\sum}_{i=1}^{n}\sum_{k\in\mathcal{Z}}\bm{1}\{u_{i}<t\}\frac{\bm{\gamma}^{k}_{i}\cdot\delta_{i}}{\mathop{\sum}\limits_{j\in\mathsf{RiskSet}(u_{i})}\mathop{\sum}\limits_{k\in\mathcal{Z}}\bm{\gamma}^{k}_{j}\bm{h}_{k}({\bm{x}}_{j},\bm{a}_{j},{\bm{k}})}$ $\hfill\triangleright$Breslow (1972)’s estimator.

$\widehat{\bm{\theta}}^{+}\leftarrow\widehat{\bm{\theta}}-\eta\cdot\nabla_{\theta}{\bm{B}}(\theta)$ $\triangleright$ Update $\bm{\theta}$ with gradient of $\widehat{Q}$.

$\widehat{\bm{\theta}}^{+}\leftarrow\mathsf{prox}_{{\epsilon}\eta}(\widehat{\bm{\theta}})$ $\triangleright$ Proximal update.
 
\returnlearnt parameters $\bm{\Omega}$;

In this section we describe the experiments conducted to benchmark the performance of SCS against alternative models for heterogenous treatment effect estimation on multiple studies including a synthetic dataset and multiple large landmark clinical trials for cardiovascular diseases.

We presented Sparse Cox Subgrouping (SCS) a latent variable approach to recover subgroups of patients that respond differentially to an intervention in the presence of censored time-to-event outcomes. As compared to alternative approaches to learning parsimonious hypotheses in such settings, our proposed model recovered hypotheses with more pronounced treatment effects which we validated on multiple studies for cardiovascular health.

While powerful in its ability to recover parsimonious subgroups there exists limitations in SCS in its current form. The model is sensitive to proportional hazards and may be ill-specified when the proportional hazards assumptions are violated as is evident in many real world clinical studies (Maron et al., 2018; Bretthauer et al., 2022). Another limitation is that SCS in its current form looks at only a single endpoint (typically death, or a composite of multiple adverse outcome). In practice however real world studies typically involve multiple end-points. We envision that extensions of SCS would allow patient subgrouping across multiple endpoints, leading to discovery of actionable sub-populations that similarly benefit from the intervention under assessment.