Identifying treatment response subgroups in observational time-to-event data

## Abstract

Abstract Identifying patient subgroups with different treatment responses is an important task to inform medical recommendations, guidelines, and the design of future clinical trials.
Existing approaches for subgroup analysis primarily rely on Randomised Controlled Trials (RCTs), in which treatment assignment is randomised.
RCTs’ patient cohorts are often constrained by cost, rendering them not representative of the heterogeneity of patients likely to receive treatment in real-world clinical practice.
When applied to observational studies, subgroup analysis approaches suffer from significant statistical biases particularly because of the non-randomisation of treatment.
Our work introduces a novel, outcome-guided method for identifying treatment response subgroups in observational studies.
Our approach assigns each patient to a subgroup associated with two time-to-event distributions: one under treatment and one under control regime.
It hence positions itself in between individualised and average treatment effect estimation.
The assumptions of our model result in a simple correction of the statistical bias from treatment non-randomisation through inverse propensity weighting.
In experiments, our approach significantly outperforms the current state-of-the-art method for outcome-guided subgroup analysis in both randomised and observational treatment regimes.

## 1 Introduction

Understanding heterogeneous therapeutic responses between patient subgroups is the core of treatment guidelines and the development of new drugs. Identifying such subgroups is valuable to inform clinical trials by identifying subgroups not responding to existing drugs, and to direct healthcare resources to those who might benefit most, and away from those who may be most harmed (Foster et al., ). For instance, subgroup analysis in the BARI trial of patients with coronary artery disease supported the use of coronary artery bypass graft over percutaneous interventions for patients with diabetes, and the opposite for patients without diabetes (Investigators, ), shaping subsequent guidelines. Figure 1 illustrates this idea: two groups may present opposite treatment responses, and would benefit from different recommendations. Our work aims to uncover such subgroups.

Randomised controlled trials (RCTs) remain the gold standard for identifying subgroups of treatment effects.
In RCTs, patients are randomly assigned to control or treated groups, allowing researchers to assess the impact of an intervention or treatment.
However, RCTs are costly and time-consuming, with restricted patient cohorts unrepresentative of the real-world diversity of patients who may receive treatment and treatment strategies (Hernán & Robins, ).

Our work introduces a novel approach that diverges from traditional RCT-based methodologies (Nagpal et al., ; ) by leveraging routinely collected observational data to identify patient subgroups. Observational studies encompass larger and more diverse cohorts reflective of real-world practices, offering the potential to uncover subgroups of treatment responses, that could be missed in RCTs. Prior works (Bica et al., ; Curth et al., ; Louizos et al., ) have leveraged observational data while addressing biases from non-random treatment assignments (Benson & Hartz, ; Hernán & Robins, ; Hernán, ). However, this body of literature primarily focuses on estimating (i) averaged treatment effects at the population level or (ii) individualised treatment effects, thereby overlooking the identification of treatment effect subgroups.

Addressing this critical gap in the literature, our work uncovers patient subgroups with distinct treatment responses using observational data. Unlike previous methods relying on RCTs, we introduce a mixture of neural networks as an extension to traditional outcome-guided models to uncover subgroups delineated by non-linear, higher-dimensional combinations of available covariates, without requiring parametrisation of the time-to-event distributions or treatment responses.

Figure: Figure 1: Subgroup treatment effect discovery in time-to-event observational data. Our method identifies subgroups of patients with similar treatment responses to guide clinical practice and design clinical trials. Our method simultaneously models the treatment effect and identifies subgroup while addressing censoring and treatment non-randomisation.
Refer to caption: extracted/5936818/images/summary.png

##### Contributions.

Our work introduces the first neural network approach to simultaneously estimate treatment effects and discover subgroups in *observational* settings, addressing both censoring and non-random treatment assignments without parameterising the survival distribution or treatment effects.
Section 2 formalises the problem of treatment subgroup discovery and non-randomisation correction. Then, we introduce the proposed monotonic neural network implementation in Section 2. Finally, we evaluate the proposed methodology on two synthetic datasets in Section 3, an extensive set of alternative settings in App. B and a real-world example described in App. C and compare it with related approaches reviewed in Section 4.

## 2 Method

### 2.1 Problem setup

##### Latent survival subgroup analysis.

Our goal is to uncover subgroups with similar treatment responses, guided by the *observed* times of occurrence of an outcome of interest.
Patients are assigned to subgroups based on their covariates at treatment time.
Each subgroup is characterised by two distributions, known as survival functions: one under treatment and one under control regimes. These distributions characterise the probability of observing the event of interest after a given time under a given treatment.

Consider the random variables associated with observed covariates $X$, an indicator identifying whether the event of interest, in our analysis death, was observed $D$, and the observed time of event $T$.
Formally, we define the latter random variables as $T:=\min(C,T^{\prime})$ and $D:=\mathbbm{1}(C>T^{\prime})$, where $C$ is random variable of the (right)-censoring time, i.e. the time at which the patient left the study before experiencing the event of interest, and $T^{\prime}$ is the partially observed random variable associated with the time of the event of interest if there was no censoring. When a patient is censored, $T=C$ and $D=0$, otherwise the event is observed and $T=T^{\prime}$ and $D=1$.
To estimate the associated survival likelihood, we further assume, as commonly done in survival analysis, that the censoring time $C$ is not informative for the time of the event of interest $T^{\prime}$.

###### Assumption 1 (Non-informative censoring) .

The censoring time $C$ is independent of the time of the event of interest $T^{\prime}$ given the covariates $X$. Formally, $T^{\prime}\perp\!\!\!\perp C\mid X$.

Central to our problem is the additional variable $Z$ associated with the latent, unobserved subgroup membership.
Following Jeanselme et al. (), we aim to identify a pre-specified $K\in\mathbb{N}$ number of subgroups(^1^11As a pre-specified number of subgroups may be a limitation in a real-world setting where we do not know the underlying grouping structure, we explore how to select this parameter based on the likelihood of the predicted outcomes in Section 3.).
We assume that $Z$ depends on $X$ and influences $T^{\prime}$. As we are interested in recovering the survival functions associated with the latent subgroups, we ignore the potential dependence of $T^{\prime}$ on $X$, considering the individual survival distribution as a mixture of the different subgroups.

###### Assumption 2 (Mixture modelling) .

The event time $T^{\prime}$ is completely determined by the patient’s group membership $Z$. Formally, $T^{\prime}\perp\!\!\!\perp X\mid Z$.

While this assumption may hurt individual performance as individual covariates do not directly inform individual survival but group membership only, it improves the interpretability of the model by separating subgroup membership from survival profiles.
In conclusion, we aim to recover $Z$ and the survival distributions associated with each subgroup from the observed $X$, $T$ and $D$.

Given the previously described dependencies, the expected survival $S$ for a given patient with covariates $x$ is the survival when marginalising over the different subgroups:

$$ $\displaystyle S(t\mid X=x):=\mathbb{P}(T^{\prime}\geq t\mid X=x)=\sum_{k=1}^{K }\mathbb{P}(Z=k\mid X=x)\mathbb{P}(T^{\prime}\geq t\mid Z=k)$ (1) $$

Note that the last term is linked to the quantity commonly used in the survival literature: $\Lambda_{k}(t):=-\log\mathbb{P}(T^{\prime}\geq t\mid Z=k)$, the group-specific cumulative hazard, with derivative with respect to time $\lambda_{k}(t)$, the group-specific instantaneous hazard, corresponding to the increase in the probability of observing the event of interest given survival until that time $t$.
Estimating the two probability distributions $\mathbb{P}(Z=k\mid X=x)$ and $\mathbb{P}(T^{\prime}\geq t\mid Z=k)$ is our core interest.

###### Assumption 1 (Non-informative censoring) .

###### Assumption 2 (Mixture modelling) .

##### Likelihood-based model fitting.

One can maximise the log-likelihood associated with the observed realisations $\{x_{i},t_{i},d_{i}\}$ of the random variables $X,T,D$ for patient $i\in[\![1,N]\!]$ with $N$ the total number of patients (we omit $i$ where redundant) to estimate the survival function (Eq. equation 1). Under Assumption 1, Equation equation 2 describes this factual(^2^22Referring to the likelihood of the observed (rather than counterfactual) realisations.) log-likelihood.

$$ $\displaystyle l_{F}(\theta)$ $\displaystyle=\sum_{i,d_{i}=1}\log\left(-\frac{\partial S_{\theta}(u\mid x_{i} )}{\partial u}\bigg{|}_{u=t_{i}}\right)+\sum_{i,d_{i}=0}\log S_{\theta}(t_{i} \mid x_{i})$ (2) $$

where $\theta$ is the set of parameters characterising the estimated survival function $S_{\theta}$.
Here, patients with an observed event ($d_{i}=1$) contribute the negative derivative of the survival function to the likelihood (see Jeanselme et al. () for derivation).
Each censored patient ($d_{i}=0$) contributes the probability of not observing an event before $t_{i}$, i.e., $S_{\theta}(t_{i}\mid x_{i})$, under assumption 1 of independence between the censoring time and the event of interest conditional on the covariates.

##### Latent treatment effect subgroups.

Consider now the binary treatment variable $A$.
Patients either receive the treatment ($A=1$) or they do not ($A=0$).
Therefore, following the potential outcomes formulation, we consider the time of events under treatment $T^{\prime}_{1}$ and under the control regime $T^{\prime}_{0}$.
The central challenge is that one can not observe both $T^{\prime}_{0}$ and $T^{\prime}_{1}$, but only $T^{\prime}$, the random variable associated with the event time under the observed treatment regime in the absence of censoring, which is equal to $A\cdot T^{\prime}_{1}+(1-A)\cdot T^{\prime}_{0}$ under assumption 4 described below. Figure 2 describes the dependencies between the previously described random variables.

A critical assumption of our proposed setting is that the subgroups we aim to identify do not influence treatment if we know a patient’s covariates, formalised as follows:

###### Assumption 3 (Unknown latent groups) .

The treatment assignment $A$ is independent of the subgroup membership $Z$ given the covariates $X$.
Formally, $A\perp\!\!\!\perp Z\mid X$.

This assumption relaxes the traditional assumption of treatment randomisation. Note that this assumption remains realistic, as (i) clinicians base treatment recommendations on the patients’ covariates, and (ii) if subgroups were known, one would not need the proposed methodology to uncover novel subgroups.

Figure: Figure 2: Graphical representation between covariates ($X$) and outcomes ($T,D$). Realisations of dashed variables are unobserved, while $X$, $A$, $T$ and $D$ are observed.
Refer to caption: extracted/5936818/images/models/NTC_DAG.png

Further, the following three assumptions are necessary to estimate treatment effects and are common in the causal literature.

###### Assumption 4 (Consistency) .

A patient’s observed event time is the potential event time associated with the observed treatment. Formally, this means $T^{\prime}:=A\cdot T^{\prime}_{1}+(1-A)\cdot T^{\prime}_{0}$ where $T^{\prime}$ is the observed event time and $(T^{\prime}_{0},T^{\prime}_{1})$ are the potential event times under the treatment $A$.

###### Assumption 5 (Ignorability) .

The potential event times are independent of the treatment given the observed covariates, i.e. $A\perp\!\!\!\perp(T^{\prime}_{0},T^{\prime}_{1})\mid X$. Equivalently, no unobserved confounders impact both treatment and event time.

###### Assumption 6 (Overlap / Positivity) .

Each patient has a non-zero probability of receiving the treatment, i.e. $\mathbb{P}(A\mid X)\in(0,1)$ where $(0,1)$ is the open interval, resulting in a non-deterministic treatment assignment.

Under the Assumptions 4, 5, and 6, existing works often focus on estimating the Individualised Treatment Effect (ITE, here denoted as $\tau$), defined as the difference between survival under the two potential treatment regimes given the covariates $x$ (see App. A.1 for derivation):

$$ $\displaystyle\tau(t,x)=\mathbb{P}(T^{\prime}\geq t\mid A=1,X=x)-\mathbb{P}(T^{ \prime}\geq t\mid A=0,X=x)$ (3) $$

Estimating this quantity requires accurately modelling the survival distributions under the two treatment regimes for all $x$ from observed data. This estimation would be straightforward if one could access the observed survival times for all patients under both treatment regimes.
In this case, one would estimate the survival using the event times under $A=1$ and under $A=0$.
The key challenge is that the counterfactual survival outcome is not observed: if a patient receives the treatment, we do not observe its outcome under no treatment, and vice versa.

In RCTs, treatment $A$ and covariates $X$ are by design independent.
Hence, the subset of the patients receiving either treatment regime is representative of the overall population.
Relying on this observed subset to estimate the survival distribution under each regime is a valid estimate for the population.
In other words, maximising the factual likelihood to estimate survival under both treatment regimes results in a valid estimate.
However, covariates and treatment are not independent in observational studies in which treatment recommendations depend upon the observed covariates. Formally, $\mathbb{P}(A\mid X)\neq\mathbb{P}(A)$ in general.
For example, clinicians might recommend more aggressive treatment to patients in more severe conditions.
This absence of randomisation results in a covariate shift between the treated and non-treated populations (Curth et al., ) as their covariate distributions differ (Bica et al., ).
In this setting, estimating survival by maximising the factual likelihood is no-longer enough to estimate the counterfactual survival distribution.

Under assumption 6, approaches such as re-weighting (Shimodaira, ), or penalisation on the dissimilarity of learnt representations (Johansson et al., ; Shi et al., ), or a combination of these approaches (Curth et al., ; Hassanpour & Greiner, ; Shalit et al., ) are remedies to estimate the counterfactual likelihood from the observed data by reducing the difference between the two treatment regimes’ populations when modelling the outcome of interest.
Specifically, Shalit et al. () demonstrate that the negative log likelihood is upper-bounded by observable quantities. Here we extend their notations to the survival setting as

$$ $\displaystyle-l(\theta)\leq-l_{F}^{*}(\theta)+\gamma\cdot\text{IPM}(q_{\Phi}^{ A=0},q_{\Phi}^{A=1}),$ (4) $$

where $l$ is the log-likelihood consisting of both factual and counterfactual log-likelihoods, $l_{F}^{*}$ is the factual log-likelihood (defined in equation 2) weighted with an inverse propensity of treatment weighting for each patient, $q_{\Phi}^{A=a}=q(\Phi(X)\mid A=a)$ is the density function associated with the transformation $\Phi$ of the covariates $X$, the Integral Probability Metric (IPM) is a distance between distributions, and $\gamma$ is a positive constant.

propose to use this bound to train a neural network to estimate treatment effects. In this context, $\Phi$ is an inner representation of the network and the IPM regularisation renders this embedding similar between the treated and untreated populations.
This embedding is then used to estimate the survival function under both treatment regimes.
The difference between the estimated survival functions is then an accurate estimate of the treatment effect as the IPM penalisation corrects for the shift resulting from treatment non-randomisation.

Note that, under the dependencies we assume, only the subgroup assignment ($\mathbb{P}(Z\mid X)$) depends on the covariates $X$ and could, consequently, play the role of the transformation $\Phi$.
However, treatment is assumed independent of the group membership under assumption 3, meaning that treatment rate does not differ between subgroups. This assumption results in the regularisation term being null, i.e. $\text{IPM}(q_{Z\mid X}^{A=0},q_{Z\mid X}^{A=1})=0$ (see App. A.2 for derivation). This assumption is consequently necessary to correctly estimate the subgroups’ treatment effect, without penalising the clustering structure.

As a consequence, our work focuses on the first term of the upper bound in equation 4 by re-weighting the factual likelihood with the patient’s propensity score (Hassanpour & Greiner, ; Shalit et al., ).
Using an estimator $\hat{p}_{A}(x):=\mathbb{P}(A=1\mid X=x)$ of the propensity score, we use a truncated propensity weighting (Austin & Stuart, ) scheme to avoid unstable weights in which the weights $w_{i}$ are defined as

$$ $\displaystyle w_{i}^{-1}=\begin{cases}0.05&\text{ if }\hat{p}_{A}(x_{i})<0.05 \\ 0.95&\text{ if }\hat{p}_{A}(x_{i})>0.95\\ a_{i}\cdot\hat{p}_{A}(x_{i})+(1-a_{i})\cdot(1-\hat{p}_{A}(x_{i}))&\text{ otherwise},\end{cases}$ (5) $$

where $i$ is the patient index,
and $a_{i}$ is the realisation of $A$ for patient $i$.
Using the factual likelihood from equation 2 and the weights $w_{i}$ defined in equation 5, we derive the upper-bound of the Negative Log-Likelihood (NLL), later used to train our model, as

$$ $\displaystyle\sum_{i,d_{i}=1}w_{i}\log\left(-\frac{\partial S_{\theta}(u\mid x _{i},a_{i})}{\partial u}\bigg{|}_{u=t_{i}}\right)+\sum_{i,d_{i}=0}w_{i}\log S_ {\theta}(t_{i}\mid x_{i},a_{i})$ (6) $$

Figure: Figure 3: Causal Survival Clustering architecture. Latent parameter $l_{k}$ characterising the subgroup $k$ is inputted in the monotonic network $M$ to estimate the cumulative hazard $\Lambda_{k}$ under both treatment regimes. $G$ assigns the probability to belong to each subgroup given the patient’s covariate(s) $x$. To tackle the challenge of treatment assignment bias, the network $W$ estimates the treatment propensity used to weigh the training likelihood. In this context, patient survival is the average of the weighted survival distributions across subgroups under the given treatment.
Refer to caption: extracted/5936818/images/models/NTC_Model.png

###### Assumption 3 (Unknown latent groups) .

###### Assumption 4 (Consistency) .

###### Assumption 5 (Ignorability) .

###### Assumption 6 (Overlap / Positivity) .

### 2.2 Estimating the quantities of interest with neural networks

The previous section discussed the quantities one must estimate —here parameterised by neural networks— to uncover subgroups of treatment effects: the assignment function $\mathbb{P}(Z\mid X)$, the cumulative incidence $\Lambda_{k}$ characterising survival under the two treatment regimes, and the weights $w$. Figure 3 illustrates the overall architecture and the neural networks used to estimate these quantities.

##### Subgroup assignment.

Similar to Jeanselme et al. (), a multi-layer perceptron $G$ with final Softmax layer assigns each patient characterised by covariates $x$ its probability of belonging to each subgroup, characterised through a $K$-dimensional vector of probabilities.

$$ $\displaystyle G(x):=[\mathbb{P}(Z=k\mid x)]_{k=1}^{K}$ $$

##### Survival distributions.

Each subgroup $k$ is represented by a vector $l_{k}\in\mathbbm{R}^{L}$ of dimension $L$, a latent parametrisation of the cumulative hazard functions. The vector $l_{k}$ is concatenated with $t$ and used as input to a neural network $M$ with monotonic positive outcomes(^3^33To ensure this constraint, we apply a square function on all weights as proposed in Jeanselme et al. (; ).), with a final SoftPlus layer to model the cumulative hazard functions under both treatment regimes $\Lambda_{k}(t):=(\Lambda_{k}(t\mid A=0),\Lambda_{k}(t\mid A=1))$.
We use the following transformation to ensure that no probability is assigned to negative times, a limitation raised concerning previous monotonic neural networks (Shchur et al., ):

$$ $\Lambda_{k}(t):=t\cdot M(l_{k},t)$ $$

By modelling these two cumulative intensity functions, we can estimate the Subgroup Average Treatment Effect, $\hat{\tau}_{k}(t)$, in subgroup $k$ as

$$ $\hat{\tau}_{k}(t):=\mathbb{E}[\tau(t,X)\mid Z=k)]=\exp(-\Lambda_{k}(t\mid A=1) )-\exp(-\Lambda_{k}(t\mid A=0))$ $$

##### Inverse propensity weighting.

Under treatment randomisation, as in RCTs, the previous components would accurately model the observed survival outcomes and identify subgroups of treatment effects by maximising the factual likelihood.
As previously discussed, to account for the treatment non-randomisation in observational studies, we weight the factual likelihood using the propensity score of a patient estimated through a multi-layer perceptron $W$ with a final sigmoid transformation as

$$ $W(x):=\mathbb{P}(A=1\mid x)$ $$

### 2.3 Training procedure

First, the network $W$ is trained to predict the binary treatment assignment by minimising the cross-entropy of receiving treatment.
Then, training all other components relies on minimising the weighted factual log-likelihood introduced in equation 6.
The use of monotonic neural networks results in the efficient and exact computation of the log-likelihood as automatic differentiation of the monotonic neural networks’ outcomes readily provides the instantaneous hazards $\lambda_{k}$ necessary for computing the survival function derivative in the log likelihood (Jeanselme et al., ; ; Rindt et al., ).

## 3 Experimental analysis

As counterfactuals are unknown in observational data, we adopt—as is common practice in this research field—a synthetic dataset in which underlying survival distributions and group structure are known.
Appendix C accompanies these synthetic results with the analysis of heterogeneity in real-world, adjuvant radiotherapy responses for patients diagnosed with breast cancer.
Our code to produce the synthetic dataset, the model, and the reproduction of all experimental results is available on Github(^4^44[https://github.com/Jeanselme/CausalNeuralSurvivalClustering](https://github.com/Jeanselme/CausalNeuralSurvivalClustering)).

### 3.1 Data generation

We generate a population of $3,000$ equally divided into $K=3$ subgroups.
We draw 10 covariates from normal distributions with different centres and survival distributions for each treatment regime following Gompertz distributions (Pollard & Valkovics, ) parameterised by group membership and individual covariates. Note that this setting relaxes the independence of outcomes and covariates assumed by the proposed model.
This choice reflects a setting in which time-to-event distributions under treatment or control regimes are different functions of the covariates. This simulation is more likely to capture the complexity of real-world responses, in contrast to traditional evaluations of subgroup analysis which often assume a linear treatment response, i.e., a shift between treated and untreated distributions.
Further, we implement two treatment assignment scenarios: a randomised assignment in which treatment is independent of patient covariates, similar to RCTs (Randomised), and one in which treatment is a function of the patient covariates as in observational studies (Observational).
Finally, non-informative censoring times are drawn following a Gompertz distribution.
Details of the generative process of our synthetic dataset are deferred to App. B.1.

### 3.2 Empirical settings

##### Benchmark methods.

We compare the proposed approach against the state-of-the-art Cox Mixtures with Heterogeneous Effect (CMHE Nagpal et al. ()(^5^55Implemented in the Auton-Survival library (Nagpal et al., ))) which uncovers treatment effect and baseline survival latent groups.
This method uses an expectation maximisation framework in which each patient is assigned to a group for which a Cox model is then fitted.
The central differences with our proposed approach is that CMHE (i) clusters patients in treatment effects subgroups and survival subgroups, and (ii) assumes a linear treatment response. This separation between survival and treatment response improves the model’s flexibility at the cost of interpretability as the number of groups grows exponentially, and subgroups of treatment effect are independent of survival. Further, the assumption of linear treatment response may hinder the discovery of subgroups with more complex responses.
By contrast, our approach identifies subgroups of treatment while considering survival, without constraining the treatment response.
We argue that these are key strengths of our method as a group not responding to treatment with low life expectancy would most benefit from alternative treatments, in comparison to a group with the same treatment responses but better survival odds. Considering jointly non-linear treatment effects and survival, therefore, results in identifying more clinically relevant subgroups.

For a fair comparison, we present three alternatives of CMHE: one with fixed $K=3$ survival subgroups, one with $L=3$ treatment effect subgroups,
and one with $K=L=2$, which allows for a total of 4 subgroups.
Crucially, these methods assume proportional hazards for each subgroup and do not account for the treatment non-randomisation.
Additionally, we also compare our model against its unadjusted alternative (CSC Unadjusted), which uses the unweighted factual likelihood ($w_{i}=1$) and should suffer from non-randomisation of treatment assignment.

Finally, we compare against two step-wise approaches in which clustering and treatment effect are modelled separately. First, we use an unsupervised clustering algorithm on the covariates, followed by a non-parametric estimate of the treatment effect as proposed in Nagpal et al. (), referred as Kmeans + TE in the following. Specifically, we use a K-Means (Hartigan & Wong, ) to cluster the data and we compute the difference between Kaplan-Meier (Kaplan & Meier, ) estimates between control and treated patients stratified by clusters. However, this approach is limited by its implicit assumption of randomisation and its pre-specified number of subgroups, which is not informed by observed outcomes. Second, we use a Virtual Twins approach to estimate individualised treatment effects and then cluster them. As proposed in the literature, we propose to use a survival tree to estimate response under each treatment regime and then use a KMeans on the estimated treatment effects, computed as the difference between survival estimates. Note that this methodology also assumes treatment randomisation, and the clustering is uninformed by the covariates.
For details on training and hyperparameter optimisation, we refer to App. B.2.

##### Evaluation.

In the synthetic experiments, the subgroup structure is known. We measure the adjusted(^6^66Random patient assignment results in an adjusted Rand-Index of 0.) Rand-Index (Rand, ), which quantifies how the estimated assignment aligns with the known underlying group structure.
Additionally, we use the integrated absolute error (IAE) between the treatment effect estimate and the ground truth, which measures how well we recover each subgroup’s treatment effect

$$ $\text{IAE}_{k}(t)=\int_{0}^{t}\left|\hat{\tau}_{\hat{k}}(s)-\tau_{k}(s)\right|ds,$ $$

where $\hat{\tau}_{\hat{k}}$ is the estimated treatment effect for subgroup $\hat{k}=\operatorname*{arg\,max}_{l}\mathbb{E}_{x\in k}(Z=l\mid x)$, i.e. the most likely assigned cluster for patients in the underlying $k$ cluster, and $\tau_{k}$ is the ground truth.

### 3.3 Treatment effect recovery

In this section, we present the results under the previously described data generation process. App. B.4 presents alternative datasets which demonstrate the flexibility of the proposed methodology.

Figure: Figure 4: Averaged negative log-likelihood across 5-fold cross-validation given the number of subgroups $K$ under the "Observational" treatment assignment with the shaded area representing 95% CI. The log-likelihood presents an elbow around the underlying number of subgroups.
Refer to caption: extracted/5936818/images/simulation/ntc_obs_loglikelihood.png

##### Recovering the underlying number of subgroups.

For all methodologies, we need to choose the number of subgroups a priori.
An important question is, therefore, whether we can identify the true number of subgroups with our model.
Figure 4 presents the average NLL obtained by cross-validation for models with different numbers of subgroups $K$.
The dotted lines represent the elbow heuristic (Thorndike, ), which identifies a change point in the explained variability of a clustering strategy, here considering the log-likelihood.
Using this heuristic, the optimal choice for $K$ is $3$, which agrees with the underlying generative process. App. B.3 explores the sensitivity to this choice, demonstrating the method’s robustness.
This data-driven choice of the number of subgroups $K$ is a crucial strength of our method, compared to classical two-stage analyses which separate clustering from treatment effect estimates. In these methods, survival outcomes cannot directly guide the choice of $K$.

##### Discovering subgroups.

Next, we present quantitative results of our method against the benchmarks methods.
Table 1 presents the performance of the different methodologies under the two studied scenarios.
Recall that $K$ denotes the number of treatment response subgroups, while in CMHE, the additional parameter $L$ describes the number of survival distributions.

Table: Table 1: 5-fold cross-validated performance averaged (with standard deviation in parenthesis) under the Randomised and Observational treatment simulation. Best performance per column and simulation scenario are marked in bold, second best in italic. ’-’ describes when the methodology diverges. Our proposed CSC method and its unadjusted variant best recovers the underlying treatment responses, with the adjusted approach presenting the best performance in the observational setting.

CSC outperforms all CMHE alternative, the current state-of-the-art method.
CMHE’s parametrisation and assumptions explain this difference. Critically, CMHE assumes (i) proportional hazards and (ii) a linear impact of treatment on log-hazards.
Neither of these two assumptions is likely to hold in real-world settings as mimicked by our synthetic data.
In contrast, our approach does not constrain the treatment effect due to its flexible modelling of the survival function under both treatment regimes.
Increasing the number of subgroups, as shown in $K=L$ improves the performance of CMHE in terms of clustering quality (Rand-Index) and the recovery of the underlying treatment effect (lower IAE), but still is inferior compared to our proposed methods. These experiments highlight the advantage of the proposed CSC method in uncovering subgroups of treatment responses due to the flexibility in modelling complex survival distributions under both treatment regimes without proportionality assumption.

CSC presents the best performance in identifying the underlying subgroup.
Our proposed CSC method outperforms all approaches on the Rand-Index. In particular, the two-step approaches do not identify the underlying subgroups when covariates and outcomes subgroups are aligned. KMeans + TE identify subgroups that are independent of the outcome of interest, as shown by the Rand-Index. Consequently, the identified clusters differ in their treatment responses as shown by the large IAE. The Virtual Twins approach presents a better capacity to identify the clusters of interest with better Rand-Index and IAE. However, the two-step approach hurts subgroup identification due to the disconnect between the clustering and input covariates. This disconnect leads to a significantly lower Rand-Index in comparison to CSC despite accurate treatment effect estimates.

Treatment assignment correction improves subgroup identification in observational settings.
The Observational simulation demonstrates the importance of correcting the likelihood under treatment non-randomisation. The two CSC alternatives present comparable performance in the randomised setting, as theoretically expected, due to a constant $w_{i}$ in this context. Note that estimating propensity slightly reduces performance due to error introduced, reducing performance for the corrected alternative. In a randomised setting, CSC Unadjusted should be preferred.
However, the superiority of CSC over its unadjusted alternative is evident in the observational setting, where CSC better recovers the different groups as shown by the Rand-Index and subgroup treatment effects.
Finally, Appendices B.4.1 to B.4.5 further validate the robustness of our method in settings with unequally distributed cluster sizes, population size, increased number of underlying clusters, different treatment rates and covariate structures.

##### Real-world heterogeneity analysis in adjuvant radiotherapy responses (App. C ).

Using the Surveillance, Epidemiology, and End Results program(^7^77Available at [https://seer.cancer.gov/](https://seer.cancer.gov/))(Seer), App. C investigates a real-world setting: we analyse which subgroups of breast cancer patients most benefit from adjuvant radiotherapy, an active area of research. Our analysis highlights a subgroup characterised by a higher number of distant lymph nodes that benefit from treatment. Despite the inherent limitations of the dataset, our exploratory method identifies a subgroup that alternative methods miss, potentially informing future clinical trials to validate or challenge this hypothesis. Critically, this demonstrates that our work offers a practical methodology for exploratory subgroup analysis of medical data to identify treatment effect heterogeneity in observational studies.

## 4 Related work

While survival subgrouping has been proposed through mixtures of distributions (Nagpal et al., ; ; Jeanselme et al., ) to identify different phenotypes of patients, the machine learning literature on phenotyping treatment effects with time-to-event outcomes remains sparse. The current literature focuses on estimating population or individual treatment effects (Curth et al., ; Shalit et al., ; Johansson et al., ; Zhang et al., ). While these models help to understand the average population-wide response to treatment and aim to estimate individualised treatment effects, they do not provide an understanding of the patient groups that may benefit more or be harmed more by treatment. Identifying such groups aligns with, and is thus more useful for, drafting medical guidelines that direct treatment to those subgroups most likely to benefit from it.

Discovering intervenable subgroups is core to medical practice, particularly identifying subgroups of treatment effect, as patients do not respond like the average (Bica et al., ; Ruberg et al., ; Sanchez et al., ) and the average may conceal differential treatment responses. Identification of subgroups has long been used to design RCTs. Indeed, subgroups identified a priori can then be tested through trials (Cook et al., ; Rothwell, ). A posteriori analyses have gained traction to uncover subgroups of patients from existing RCTs to understand the underlying variability of responses.

The first set of subgroup analysis methodologies consists of a step-wise approach: (i) estimate the ITE and (ii) uncover subgroups using a second model to explain the heterogeneity in ITE. Foster et al. (); Qi et al. () describe the virtual twins approach in which one models the outcome using a decision tree for each treatment group. The difference between these decision trees results in the estimated treatment effects. A final decision tree aims to explain these estimated treatment effects to uncover subgroups. Similar approaches have been explored with different meta-learners (Xu et al., ), or Bayesian additive models (Hu et al., ), or replacing the final step with a linear predictor to uncover the feature influencing heterogeneity (Chernozhukov et al., ). However, Guelman et al. () discuss the drawbacks associated with these approaches. Notably, the two-step optimisation may not lead to recovery of the underlying subgroups of treatment effects.

Tree-based approaches were proposed to address the limitations of step-wise approaches by jointly discovering subgroups and modelling the treatment effect. Instead of traditional splits on the observed outcomes, these causal trees aim to discover homogeneous splits regarding covariates and treatment effects. Su et al. () introduce a recursive population splitting based on the average difference in treatment effect between splits. Athey & Imbens (; ) improve the confidence interval estimation through the honest splitting criterion, which dissociates the splitting from the treatment effect estimation. Wager & Athey () agglomerate these causal trees into causal forests for improved ITE estimation. Each obtained split in the decision tree delineates two subgroups of treatment effect Lipkovich et al. (); Loh et al. (). Alternatively, McFowland III et al. () propose pattern detection and Wang & Rudin (), causal rule set learning to uncover these subgroups. However, all these approaches rely on a local optimisation criterion (Lipkovich & Dmitrienko, ) and greedy split exploration. Recently, Nagpal et al. () addressed the local optimisation by constraining the treatment response to a linear form in a mixture of Cox models.

Previous approaches uncover subgroups of treatment effect but consider RCTs with binary outcomes, not the observational setting with inherently continuous survival outcomes that our paper explores. At the intersection with survival analysis, Zhang et al. () extend causal trees to survival causal trees, modifying the splitting criterion by measuring the difference in survival estimates between resulting leaves. Similarly, Hu et al. () propose Bayesian additive models and Zhu & Gallego () propose a step-wise approach with propensity weighting to study observational data. Closest to our work, (Jia et al., ; Nagpal et al., ) propose to uncover subgroups within RCTs with survival outcomes. Jia et al. () propose a mixture of treatment effects characterised by Weibull distributions trained in an expectation-maximisation framework. Similarly, Nagpal et al. () stratify the population into three groups: non-, positive- and negative responders to treatment. An iterative Monte Carlo optimisation is used to uncover these subgroups, characterised by a Cox model with a multiplicative treatment effect. As demonstrated in our work, this step-wise optimisation may be limiting, and the assumption of RCTs renders the model less relevant in observational data.

## 5 Conclusion

This work fills the current gap in the literature of methodologies to identify subgroups of treatment responses in observational time-to-event data.
While modelling the observed outcome, the approach jointly identifies patient subgroups with different treatment responses. Our experiments demonstrate the capacity of our proposed method to uncover groups of treatment effects in both synthetic and real-world datasets. However, causal modelling relies on empirically unverifiable assumptions. The invalidity of these assumptions is a lesser concern in our targeted setting of hypothesis generation, in comparison to individual estimation of treatment recommendations.

Serving the purpose of a hypothesis-generating tool, we invite practitioners to further investigate the subgroups identified in observational studies through RCTs to (i) validate the estimated responses, (ii) identify potential alternative treatments with improved responses, and (iii) inform clinical guidelines.

#### Acknowledgements

The authors acknowledge the partial support of the UKRI Medical Research Council (MC_UU_00002/5 and MC_UU_00002/2).
VJ, CHY and FF acknowledge the Enrichment Scheme of The Alan Turing Institute under the EPSRC Grant EP/N510129/1.
FF acknowledges the receipt of studentship awards from the Health Data Research UK-The Alan Turing Institute Wellcome PhD Programme (Grant Ref: 218529/Z/19/Z).

## Appendix A Proof

### A.1 Individualised Treatment Effect

This section derives the individualised treatment effect expression introduced in equation 3.

$$ $\displaystyle\tau(t,x):$ $\displaystyle=\mathbb{E}(\mathbbm{1}(T_{1}\geq t)-\mathbbm{1}(T_{0}\geq t)\mid X =x)$ $\displaystyle=\mathbb{E}(\mathbbm{1}(T_{1}\geq t)\mid X=x)-\mathbb{E}(\mathbbm {1}(T_{0}\geq t)\mid X=x)$ $\displaystyle=\mathbb{P}(T^{\prime}\geq t\mid A=1,X=x)-\mathbb{P}(T^{\prime} \geq t\mid A=0,X=x)$ (Under assumption 4 and 5 ) $$

### A.2 Upper-Bound Simplification

In Section 2, we claim that under the considered DAG presented in Figure 2,

$$ $\text{IPM}(q_{Z\mid X}^{A=0},q_{Z\mid X}^{A=1})=0$ $$

with $A$, being the treatment assignment. Under assumption 3, expressing this quality corresponds to:

$$ $\displaystyle q_{Z\mid X}^{A=a}:=q(Z\mid X,A=a)=q(Z\mid X)$ $$

with $q$, the density function. From this expression, $q_{Z\mid X}^{A=0}=q_{Z\mid X}^{A=1}$, which results in the distance between these distributions being null.

## Appendix B Synthetic analysis

### B.1 Data generation

We consider a synthetic population of $N=3,000$ patients with 10 associated covariates $X\in\mathbb{R}^{10}$ divided into $K=3$ subgroups. The following data generation does not aim to mimic a particular real-world setting but follows a similar approach to Nagpal et al. (). The following describes our generation process:

##### Covariates.

Each patient’s membership $Z$ is drawn from a multinomial with equal probability. Group membership informs the two first covariates through the parametrisation of the bivariate normal distribution with centres $c_{k}$ equal to $(0,2.25)$, $(-2.25,-1)$, and $(2.25,-1)$. All other covariates are drawn from standard normal distributions. Formally, this procedure is described as:

$$ $\displaystyle Z$ $\displaystyle\sim\text{Mult}\left(1,\left[\frac{1}{3},\frac{1}{3},\frac{1}{3} \right]\right)$ $\displaystyle X_{[1,2]}\mid Z=k$ $\displaystyle\sim\text{MVN}(c_{k},I^{2})$ $\displaystyle Y$ $\displaystyle\sim\text{Mult}\left(1,\left[\frac{1}{4},\frac{1}{4},\frac{1}{4}, \frac{1}{4}\right]\right)$ $\displaystyle X_{[3:10]}\mid Y=k$ $\displaystyle\sim\text{MVN}(c^{\prime}_{k},I^{8})$ $$

with MVN denoting a multivariate normal distribution, $c^{\prime}_{k}$ random cluster centres, and $I^{n}$, the identity covariance matrix of dimension $n$. Note that we introduce $Y$ to present a covariates structure that is independent of the treatment responses.

##### Treatment response.

For each subgroup, event times under treatment and control regimes are drawn from Gompertz distributions, with parameters that are functions of group-specific coefficients ($B^{0}$ and $\Gamma^{0}$ for the event time when untreated and $B^{1}$ and $\Gamma^{1}$ when treated) and the patient’s covariates.

$$ $\displaystyle B^{0}_{z}\mid Z=z$ $\displaystyle\sim\text{MVN}(0,I^{10})$ $\displaystyle\Gamma^{0}_{z}\mid Z=z$ $\displaystyle\sim\text{MVN}(0,I^{10})$ $\displaystyle T_{0}\mid Z,X,B^{0}_{z},\Gamma^{0}_{z}=(z,x,\beta_{z}^{0},\gamma _{z}^{0})$ $\displaystyle\sim\text{Gompertz}\left(w_{0}(\beta_{z}^{0},x),s_{0}(\gamma_{z}^ {0},x)\right)$ $\displaystyle B^{1}_{z}\mid Z=z$ $\displaystyle\sim\text{MVN}(0,I^{10})$ $\displaystyle\Gamma^{1}_{z}\mid Z=z$ $\displaystyle\sim\text{MVN}(0,I^{10})$ $\displaystyle T_{1}\mid Z,X,B^{1}_{z},\Gamma^{1}_{z}=(z,x,\beta_{z}^{1},\gamma _{z}^{1})$ $\displaystyle\sim\text{Gompertz}\left(w_{1}(\beta_{z}^{1},x),s_{1}(\gamma_{z}^ {1},x)\right)$ $$

with $w_{0}$, $w_{1}$ two functions parametrising the Gompertz distributions’ shape as $w_{0}(\beta,x):=|\beta[0]|+(x[5:10]\cdot\beta[5:10])^{2}$, $w_{1}(\beta,x):=|\beta[0]|+(x[1:5]\cdot\beta[1:5])^{2}$, and the shift parameter parameterised as $s_{0}(\gamma,x):=|\gamma[0]|+|(x[1:5]\cdot\gamma[1:5])|$ and $s_{1}(\gamma,x):=|\gamma[0]|+|(x[5:10]\cdot\gamma[5:10])|$ where $v[i]$ described the $i^{th}$ element of the vector $v$. These functions aim to introduce non-linear responses with discrepancies between control and treatment regimes. Note that we allow covariates to influence the survival distribution as a patient’s covariates influence Gompertz’s shapes and scales.

##### Treatment assignment.

The non-randomisation of treatment is central to the problem of identifying treatment subgroups in real-world applications. Assuming a treatment assignment probability of 50%, we assign each patient to a given treatment. We propose two treatment assignment strategies reflecting a RCT and an observational setting, denoted as "Randomised" and "Observational". "Randomised" consists of a Bernoulli draw using the realisation of $P$. "Observational" reflects an assignment dependent upon the observed covariates.

$$ $\displaystyle A_{rand}\mid P=p$ $\displaystyle\sim\text{Bernoulli}(0.5)$ $\displaystyle A_{obs}\mid P,X=(p,x)$ $\displaystyle\sim\text{Bernoulli}(F_{\Phi(X)}(\Phi(x))\times 0.5)$ $$

with $F_{\Phi(X)}(\Phi(x))$ the cumulative distribution function that returns the probability that a realisation of a $\Phi(X)$ will take a smaller value than $\Phi(x)$. In our experiment, we chose $\Phi(x)=\sum x^{2}$.

##### Censoring.

Finally, our work focuses on right-censored data. To generate censoring independent of the treatment and event, we draw censoring from another Gompertz distribution as follows:

$$ $\displaystyle B^{C}$ $\displaystyle\sim\text{MVN}(0,I^{5})$ $\displaystyle C\mid X,B^{C}=(x,\beta)$ $\displaystyle\sim\text{Gompertz}\left(w_{c}(\beta,x),0\right)$ $\displaystyle T^{\prime}$ $\displaystyle=A\cdot T_{1}+(1-A)\cdot T_{0}$ $\displaystyle T$ $\displaystyle=\min(C,T^{\prime})$ $\displaystyle D$ $\displaystyle=\mathbbm{1}(C>T^{\prime})$ $$

with $w_{c}:=(x[5:10]\cdot\beta)^{2}$, the scale of the censoring Gompertz distribution.

From these data, the goal is to model the treatment effect from the observed $X,T,D,A$ with $T\mid A=a,T_{0},T_{1},C=t_{0},t_{1},c:=\min(c,(1-a)*t_{0}+a*t_{1})$ and $D=\mathbbm{1}_{C>T}$.

### B.2 Training and hyperparameter optimisation

##### Training.

We perform a 5-fold cross-validation for both Randomised and Observational simulations.
For each cross-validation split, the development set is divided into three parts: 80% for training, 10% for early stopping, and 10% for hyper-parameter search over a grid presented in Appendix B.2.
All models were optimised for 1000 epochs using an Adam optimiser (Kingma & Ba, ) with early stopping.

##### Hyperparameter optimisation.

We adopted a 100-iteration random grid-search over the following hyperparameters: network depth between 1 and 3 inner layers with 50 nodes, latent subgroup representation in $[10,25,50]$ and for training, a learning rate of $0.001$ or $0.0001$ with batch size of $100$ or $250$.

### B.3 Sensitivity

In Section 3, we describe how the proposed methodology presents an elbow in the likelihood as a function of the number of clusters around the underlying number of clusters $K=3$. This section explores how the identified treatment effects change under a misspecified model which estimates 2 and 4 clusters. This analysis examines the risk associated with misspecifying the number of clusters.

Figure 5 illustrates the cross-validated subgroup treatment effects when trained with $K=2,3$ and $4$ clusters. Uncertainty increases with misspecified models due to the instability of the identified subgroups.

Figure: (a) $K=2$
Refer to caption: extracted/5936818/images/simulation/k=2.png

### B.4 Alternative data generations

In this section, we explore alternative data generations to demonstrate the robustness of the proposed strategy.

#### B.4.1 Unequal cluster size

Section B.1 presents a population equally distributed over the different clusters. In medical settings, the underlying subgroups may differ in size. This section explores an alternative scenario in which the population is distributed over the three clusters: 62.5%, 25% and 12.5%.

Table 2 summarises the performance in this simulation, echoing the main text’s conclusions. The proposed method best recovers the different subgroups and presents one of the best estimates of subgroups’ treatment effects..

Table: Table 2: Cross-validated performance (with standard deviation in parenthesis) when clusters are of different sizes.

#### B.4.2 Number of points

Section B.1 describes a data generation with $3,000$ patients. This section presents results when there are $300$ and $30,000$ patients. A smaller number of patients may result in a larger impact of treatment non-randomisation but also impact the neural network capacity to identify underlying subgroups of treatment effect. A larger number makes non-randomisation less of a concern.

Table 3 presents the performance with these different population sizes under an observational treatment assignment. A first observation is that all methodologies present better performance with larger number of points. Further, baselines that do not account for the assignment mechanisms present lower performance with $N=300$ as non-randomisation has an increased impact on treatment effect estimate with a smaller population. This difference decreases when $N=30,000$ with the virtual twins approach presenting the best recovery of the treatment effects. However, throughout the different settings, the proposed CSC presents the best Rand-Index indicating a good recovery of the underlying structure.

Table: Table 3: Cross-validated performance (with standard deviation in parenthesis) with varying $N$ under observational treatment setting.

#### B.4.3 Number of clusters

Section B.1 describes a data generation with $K=3$ clusters. This section explores the behaviour of the different methodologies with an increased number of clusters $K=5$. Following the same data generation, we introduces two additional clusters with centres [-3, 3] and [4, 4].

Table 4 presents the different models’ performance in this setting, echoing the same results as the main text. The proposed method best recovers the underlying clustering structure as shown by the Rand-Index. CMHE suffers from its inherent association between treatment and survival. The step-wise approach is unable to identify the clusters of interest if the data presents a structure that is not fully aligned with the outcome of interest.

This scenario further validates our findings, demonstrating the methodology’s capacity to generalise beyond the scenario presented in the main text.

Table: Table 4: Cross-validated performance (with standard deviation in parenthesis) when $K=5$.

#### B.4.4 Impact of treatment rate

Section B.1 assumes 50% of the population receives treatment. This section explores when 25% and 75% of the population receive treatment under the non-randomised treatment setting.

Table 5 presents the performance under these different treatment rates in an observational setting. These results highlight CSC’s capacity to identify the underlying subgroups with the highest Rand-Index and one of the best cluster treatment effect recovery in these settings.

Table: Table 5: Cross-validated performance (with standard deviation in parenthesis) with varying treatment rates.

#### B.4.5 Aligned covariates and treatment effect structure

Section B.1 presents a data generation with the last 8 covariates having a clustered structures independent of the treatment response. This section explores a setting where all clusters in the covariates are associated with different treatment responses of interest. Specifically, we sample the last dimensions from standard normal distributions instead of various clusters, i.e.,

$$ $X_{[3-10]}\sim\text{MVN}(0,I^{8})$ $$

Table 6 summarises the performance in this setting, evidencing an improvement of the Kmeans+TE approach. When covariates’ clusters are aligned with the outcome of interest, this method recovers the clustering structure well, as shown by the best Rand-Index. Non-randomisation of treatment, however, worsens the treatment estimates due to the absence of correction in the estimation of treatment responses.

Even in this unrealistic scenario, our proposed method remains the second best in recovering the underlying clustering structure and associated treatment responses. This observation reinforces the main findings of our work, demonstrating the method’s potential to identify subgroups of treatment effects across settings.

Table: Table 6: Cross-validated performance (with standard deviation in parenthesis) when all covariates clusters presents different treatment responses. This setting aligns with the assumptions made by the step-wise approaches.

#### B.4.6 Linear treatment responses

As an alternative setting, we propose a setting with a linear treatment response, as assumed by CMHE Nagpal et al. (). To this end, we followed a similar process as described in Section B.1 at the exception of the treatment response. We draw the response under no treatment from a Gompertz, and the one under treatment with the same parameters shifted by the treatment effect. Formally, using the same notation

$$ $\displaystyle B^{0}_{z}\mid Z=z$ $\displaystyle\sim\text{MVN}(0,I^{10})$ $\displaystyle\Gamma^{0}_{z}\mid Z=z$ $\displaystyle\sim\text{MVN}(0,I^{10})$ $\displaystyle T_{0}\mid Z,X,B^{0}_{z},\Gamma^{0}_{z}=(z,x,\beta_{z}^{0},\gamma _{z}^{0})$ $\displaystyle\sim\text{Gompertz}\left(w_{0}(\beta_{z}^{0},x),s_{0}(\gamma_{z}^ {0},x)\right)$ $\displaystyle T_{1}\mid Z,X,B^{1}_{z},\Gamma^{0}_{z},A_{z}=(z,x,\beta_{z}^{0}, \gamma_{z}^{0},a_{z})$ $\displaystyle\sim\text{Gompertz}\left(w_{0}(\beta_{z}^{0},x)\times a_{z},s_{0} (\gamma_{z}^{0},x)\right)$ $$

with $a_{z}$ the treatment effect for cluster $z$. In this experiments, we choose $a_{0},a_{1},a_{2}=(0.5,1,1.5)$.

Table 7 summarises performances in this scenario. While demonstrating a significant improvement of the CMHE alternatives, CSC presents competing performance on subgroup identification and treatment effect estimates.

Table: Table 7: Cross-validated performance (with standard deviation in parenthesis) when considering a linear treatment response. This setting aligns with the assumption made by CMHE.

## Appendix C Real-world case Study: Heterogeneity of adjuvant radiotherapy responses in the Seer dataset

To study the medical relevance of the proposed method, we explore data from the Surveillance, Epidemiology, and End Results program(^8^88Available at [https://seer.cancer.gov/](https://seer.cancer.gov/)) (Seer) gathering patients diagnosed with breast cancer between 1992 and 2017. Following Lee et al. (); Danks & Yau (); Jeanselme et al. (), we select women who died from the condition or from cardiovascular diseases. From this observational data, we investigate the impact of adjuvant radiotherapy after chemotherapy on survival outcomes. To this end, we subselect patients with recorded treatment and who received chemotherapy. These criteria led to the selection of 239,855 patients with 22 covariates measured at diagnosis such as diagnosis year, grades, ethnicity, laterality, tumour size and type (see  Danks & Yau () for further description).

##### Selection of K 𝐾 K italic_K .

From the selected population, our aim is to identify heterogeneous responses to adjuvant treatment. The first challenge is the selection of the number of groups to use ($K$). We advise to follow medical actionability and consider the change in treatment effects and size of the subgroups when increasing this parameter. In the absence of experts’ intuition, one may rely on an elbow rule heuristic over $l_{F}$. Using Figure 6, the negative log-likelihood presents an elbow for $K=2$.

Figure: Figure 6: Cross-validated negative log-likelihood as a function of the number of groups ($K$).
Refer to caption: extracted/5936818/images/seer/ntc_seer_loglikelihood.png

##### Treatment effect subgroups.

In this section, we examine the treatment response following adjuvant radiotherapy to identify groups of patients who received surgery and chemotherapy and may benefit from adjuvant radiotherapy. This problem is central to patients’ treatment as no evidence-based guidelines for adjuvant therapy exist Lazzari et al. (), making this setting more likely to meet the positivity assumption (assumption 6), necessary to study causality in observational data.

Figure: (a) Neural Survival Treatment
Refer to caption: extracted/5936818/images/seer/ntc_seer_clusters.png

##### Uncovering treatment response.

Figure 7 presents the identified treatment effect subgroups when using CSC and CMHE, with the number of subgroups $K$ selected through hyperparameter tuning. As shown by Figure 6, using the previously described elbow rule leads to the same number of subgroups. As previously mentioned, our proposed methodology presents two strengths that explain the difference in the identified subgroups of treatment effects. First, the survival distribution under treatment is not constrained by the one under the control regime, resulting in more flexible, non-proportional distributions. CMHE’s parametrisation, which characterises treatment as a linear shift in the log hazard, results in a proportionality assumption between treated and untreated distributions. Second, CMHE does not account for treatment non-randomisation in its average treatment effect estimation, whereas our use of inverse propensity weighting corrects for any observed ones.

Using a permutation test, we identify as the covariates that most impact the likelihood associated with the model. Figure 8 displays the 10 covariates most indicative of the different treatment response subgroups.

Figure: Figure 8: Causal Survival Clustering - Change in log-likelihood given random permutation of a given covariates.
Refer to caption: extracted/5936818/images/seer/ntc_seer_importance.png

Table 8 summarises the average value across the identified subgroups and the life expectancy gain when using adjuvant radiations measured through the Restricted Mean Survival Time (RMST) Royston & Parmar (). Both methodologies identify a population with limited treatment response. However, our proposed methodology identifies a second group, characterised by larger HER2 and larger distant lymph node count, with a positive treatment response, gaining more than half a year of life expectancy over the five years following diagnosis.

Table: Table 8: Causal Survival Clustering subgroups’ characteristics in the Seer cohort described through percentage / mean (std).

The proposed analysis pinpoints a group that could benefit from adjuvant radiotherapy. However, our methodologies remain hypothesis-generating tools, requiring further experimental validation, particularly due to potential confounding through hormonal therapy (not available in this dataset), the temporal nature of treatment, and the plurality of treatment options. The quality of available covariates limits this analysis and only serves as an example to medical practitioners to identify subgroups of treatment responses from observational data.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: Line 155 - I do not think treatment randomization is a "traditional" assumption in the context of observational studies?
Do you perform hyperparameter optimization for the baselines, and if so, where are they described?
soundness: 3
strengths: Heterogeneity of treatment effects is an important problem.
Authors study this problem in the context of observational studies and under censoring.
The work is well motivated and positioned in the relevant literature.
Although not very novel component-wise, the end-to-end model architecture may be valuable for future work.
summary: the authors study the problem of identifying subgroups that respond to a treatment differently. they consider a setting with observational study and where the outcomes are right censored. they propose using ML models to flexibly model the nuisance function, and as far as I can tell, this is the main contribution in the paper.
weaknesses: My main criticism toward the work is that it is too assumption heavy. Authors assume observational study is perfect for causal analysis (e.g., no unmeasured confounding) and that the censoring time is conditionally independent of the time-to-event variable. On top of those, the assumption that the number of subgroups is known a priori makes the setting too simple/unrealistic. 

While this may not necessarily be a reason for criticism, the results derived in this work follow rather straighforwardly via an amalgamation of the existing work in the literature, bringing the amount of non-incremental contribution made by this work under question, which for me the key takeaway being using nonlinear models for the nuisance functions.

Given this restrictive assumptions, I think it is valuable to evaluate how well the models perform under different violations of the assumptions. The authors could run ablation studies, isolating the violations of different assumptions one at a time to analyze the sensitivity to those assumptions (e.g., no unmeasured confounding and independent censoring and time-to-event). 

Authors seem to do this for choosing the number of clusters in Appendix B.3., and to me, there seems to be a non-trivial degrade in performance (larger confidence intervals). This should be more clearly acknowledged in the main text and supported with more discussions in the Appendix for the underlying reasons. In the current form, it seems to be brushed off by saying the model is robust to this, but I do not totally agree.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: -	Since the assumptions seem pretty hard, how robust is the method against possible violations such as of Assumptions 2 and 3?
-	What exactly is the benefit of identifying subgroups compared to normal CATE estimation if they are not interpretable? E.g., in line 1306 a permutation test is used to interpret the subgroups, should such interpretability not be directly incorporated in to the objective?
-	For further questions see weaknesses above.
soundness: 2
strengths: -	The paper addresses an interesting topic in causal inference of clinical relevance going beyond classical CATE estimation to handle subgroup identification and time to event data from RWD / observational data.
-	The combination of the building blocks is comprehensibly explained and mostly well-motivated
-	The authors provide code for reproducibility
summary: This paper proposes a new method for identifying subgroups in treatment response estimation for observation time-to-event data. For that, the authors combine ideas from treatment effect estimation adjusting for selection bias, clustering and survival analysis. They further provide experiments on synthetic data and perform an analysis of real-world data.
weaknesses: -	The structure of the paper could be clearly improved to clarify the exact research gap, comparison to existent works, and contribution:
    -	In the introduction (line 44-52), method section (l. 201-255), and related work (l.472-479), this paper references mostly CATE methods in the static setting and thus neglecting to explain the connection to CATE methods on survival data (e.g., SurvITE, etc.) or subgroup identification for CATE which makes it hard to judge the novelty of the work.
    -	The potential outcome framework could be introduced more formally and earlier to properly describe the causal nature of the problem
    -	Instead of normal survival forests with clustering, the causal survival forest could be used as a baseline, e.g. also with the R-learner / DoubleML to adjust of selection bias.
    -	The motivation for the subgrouping could be improved, e.g., usually the objective is to get interpretable subgroups
    -	In Eq. (5) and (6), the authors should clearly explain how and why inverse propensity weighting benefits their method, e.g., by properly showing the theoretical benefits under selection bias.
-	The Equations and theoretical motivations seem to be inaccurate and could be improved and extended:
    -	Eq. (4): what exactly is $l_F^*$, what exactly does the “log-likelihood consisting of both factual and counterfactual[...]” represent? Especially wrt to the motivation of IPW later, this is an important part.
    -	What are the consequences and motivation of propensity clipping in Eq.(5) “avoid unstable weights”? (introducing bias but lowering variance, still unbiased but faster convergence?)
    -	Eq. (6) should present an equation or objective and could be better motivated
    -	$u$ should be introduced in Eq. (2) & (6)
- The contribution seems to be mainly incremental, combining existing ideas from subgroup analysis, clustering for survival analysis, and inverse propensity weighting

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: See the weaknesses. 

1. Figure 3, line 216: The proposed method uses a separate network W to estimate the propensity score. Could a simple logistic regression suffice instead? I think it’s not always necessary to use complex models—sometimes, simpler approaches work just as well.
2. Figure 7, line 1292: In subfigure (b), there appears to be only one cluster. Is this an error?
soundness: 2
strengths: 1. Unlike traditional methods that often assume linear treatment responses, the proposed model employs neural networks to allow for non-linear survival functions, making it more adaptable to real-world applications.
2. The integration of monotonic neural networks with IPW is both technically advanced and practically relevant for observational study analysis. This approach helps correct non-random assignment bias without requiring parameterization of the underlying survival distributions.
summary: The paper introduces a novel neural network-based method for identifying patient subgroups with distinct treatment responses using observational data. This approach tackles the challenges of non-random treatment assignment and censoring, which often introduce biases in observational studies, by incorporating inverse propensity weighting and a monotonic neural network to model complex, non-linear survival distributions under different treatment regimes. Empirical results are presented on both synthetic and real-world datasets.
weaknesses: 1. The main limitation lies in the interpretability of subgroups, especially regarding how these identified subgroups can be understood and applied by clinical practitioners. For example, in Appendix C, it would be helpful to know the characteristics of the two subgroups. If a patient presents with specific features, how would doctors classify them into a subgroup? Additionally, how statistically stable are these subgroups?
2. The presentation of the paper is somewhat tedious. The problem setup section (lines 89 to 289) contains an excess of basic concepts that are likely familiar to scholars in the field, making it less engaging. I recommend condensing the key assumptions and focusing on the unique aspects of the proposed approach.
3. In the synthetic experiment, the number of clusters is known to be three, and the authors use the negative log-likelihood to identify the optimal number of clusters. This heuristic indeed suggests three clusters, aligning with the generative process of the synthetic data. However, in the appendix for the real-world experiment, the cluster number is set to two, which seems questionable, with no verification provided for its accuracy. In summary, the proposed method assumes a pre-determined cluster number, which itself is a complex issue. I am concerned that the performance of the new method may heavily rely on the correctness of this cluster number, and the paper does not provide a robust solution to address this dependency.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: No.
soundness: 2
strengths: - The work focuses on model development on the observational data which is vital in this field as clinical trial data are expensive and limited. The study has considered removing potential biases in the data.
- The study models on the survival data, which are prevalent in the medical field, and the proposed method also consider the censoring event. 
- The study evaluates the proposed methods in terms of accuracy in subgroup identification and treatment effect estimation, which is solid.
- The study has demonstrated thorough ablation studies, in terms of cluster size, sample size, treatment rate.
summary: The study proposes a model to identify subgroups in terms of treatment effects based on observational time-to-event data. The major contribution of this paper is that it is the first work to combine subgroup analysis, treatment effects analysis, and survival analysis together.
weaknesses: - There was one famous work called "Deepsurv"[1] in this field which was published in 2018. Although it does not consider subgroup identification, but I think you should compare your work with theirs in terms of IAE. 
- For identifying the subgroup analysis, the work states that it utilizes a vector $l_k$ (please consider change the notation to $v$ as in your previous sections $l$ is used for log-liklihood), however, there is no description how you got the $l_k$. Is $l_k$ a trainable parameters for your model to learn or some features extracted directly from the baseline table? I would give you a borderline score for now as I didn't get your point at this time but will consider increasing the score if you can explain it clearly in the rebuttal.
- As illustrated in this paper[2], another key information when estimating the treatment effects over time-to-event data is the confidence interval to quantify the how much difference do you observe in the identified subgroup. I think it does matter in the medical science. 
- Lastly, I would say the contribution of this work is limited in both methods and application scenario. The methods uses a neural network to model the survival distrbution and estimate the treatment effects is not something new as the work [1] has proposed simiar idea in 2018. The work's framework for addressing the potential bias in observational data is also a classic stragegy used in this field as earliest proposed in "DragonNet"[3]. And in terms of the application scenario, I didn't see any advantage of your method over other external interpretation method as you illustrated in your real-world case study. In your appendix Figure 8, my understanding is that it is very similar to Shapley value like Figure 5 in [2], and it seems that it does not consider the interaction effects based on your interpretation. 

[1] Katzman, J. L., Shaham, U., Cloninger, A., Bates, J., Jiang, T., & Kluger, Y. (2018). DeepSurv: personalized treatment recommender system using a Cox proportional hazards deep neural network. BMC medical research methodology, 18, 1-12.

[2] Jiang, L., Xu, C., Bai, Y., Liu, A., Gong, Y., Wang, Y. P., & Deng, H. W. (2024). Autosurv: interpretable deep learning framework for cancer survival analysis incorporating clinical and multi-omics data. NPJ precision oncology, 8(1), 4.

[3] Narendra, K. S., & Mukhopadhyay, S. (1997). Adaptive control using neural networks and approximate models. IEEE Transactions on neural networks, 8(3), 475-485.