Integrating Copula-Based Random Forest and Deep Learning Approaches for Analyzing Heterogeneous Treatment Effects in Survival Analysis

Abstract

This paper presents deep learning models—specifically, Long Short-Term Memory (LSTM) networks and hybrid Convolutional Neural Network–LSTM (CNN-LSTM) with a Copula-Based Random Forest (CBRF) model to estimate Heterogeneous Treatment Effects (HTEs) in survival analysis. The proposed method is designed to capture non-linear relationships and temporal dependencies in clinical and genomic data, with a particular focus on exploring how treatment effects vary by race as a moderating factor. Using breast cancer data from the TCGA-BRCA dataset, which includes both clinical variables and gene expression profiles, we filter the data to focus on two racial groups: Black or African American and White. Dimensionality reduction is performed using Principal Component Analysis (PCA). We compare the CNN-LSTM, LSTM, and CBRF models under three weighting strategies—no weights, Horvitz–Thompson (HT) weights, and Inverse Probability of Treatment Weighting (IPTW)—for predicting treatment effects. Model performance is evaluated using Root Mean Squared Error (RMSE), Mean Absolute Error (MAE), Concordance statistic (C-statistic), Average Treatment Effect (ATE), and Conditional Average Treatment Effect (CATE) by race. The CNN-LSTM model consistently outperforms the others, achieving the lowest prediction errors and highest discrimination, particularly under IPTW. Among the weighting strategies, IPTW yields the most substantial improvements in model performance and bias reduction. Importantly, race-specific treatment effects exhibit notable variation: CNN-LSTM estimates a slightly higher CATE for Black individuals under IPTW. Overall, CNN-LSTM with IPTW is recommended for robust and equitable causal inference, especially in racially stratified settings.

1. Introduction

Over the past five years, researchers have significantly advanced the estimation and application of heterogeneous treatment effects (HTEs). These advancements result from both methodological innovations and an increasing number of applications across diverse fields, particularly in clinical and healthcare domains. Recent studies integrate advanced machine learning techniques, data fusion strategies, and statistical models to improve the precision of treatment effect estimation [1]. These developments have important implications for clinical practice, especially in enabling personalized treatment strategies. Looking ahead, the continued integration of machine learning, Bayesian modeling, and meta-analytic frameworks promises to enhance our understanding and application of HTE across various healthcare settings [2,3]. For example,  [2] apply flexible causal machine learning estimators in a comprehensive case study that explores the heterogeneous survival effects of two radiotherapy strategies for localized high-risk prostate cancer. Understanding how treatments affect distinct subgroups, especially those defined by race, gender, or other demographic variables remains a major challenge in medical research. Analyzing HTE across such subpopulations deepens our understanding of treatment efficacy and supports the design of more personalized and equitable healthcare interventions [4,5,6].

HTE plays a critical role in survival analysis because treatment effects rarely remain uniform across all individuals [7]. Patients differ in baseline characteristics, risk factors, comorbidities, and biological responses to interventions. Overlooking this heterogeneity may lead to inaccurate conclusions about treatment efficacy, resulting in suboptimal or potentially harmful decisions for specific subgroups [8]. Accurately modeling HTE enables the development of more personalized treatment strategies, better allocation of healthcare resources, and improved patient outcomes [9]. Nevertheless, existing methods for estimating HTE in survival contexts exhibit several limitations. Traditional models like the Cox proportional hazards model often assume constant treatment effects across populations, potentially masking significant subgroup differences [10]. Parametric and semiparametric models frequently impose rigid assumptions, such as proportional hazards, which may not hold in complex, real-world data settings [11]. While machine learning methods adapted for HTE, including causal forests, have been extended to handle survival data [12], they still struggle to integrate censoring appropriately without strong assumptions. Ref. [13] introduce DeepSurv, a deep neural network variant of the Cox model that captures interactions between patient covariates and treatment effects to generate personalized recommendations. To evaluate recent survival machine learning methods, such as DeepSurv,  [2] conducted a simulation study across varied settings involving confounded HTE and covariate overlap. Ref. [14] propose a deep learning architecture for dynamic survival modeling using longitudinal data.

Copula-based models have recently emerged as a promising approach to address these challenges [15]. By capturing the dependence structure between random variables, copulas enable more flexible modeling of relationships between covariates and treatment effects [16]. These methods capture nonlinear dependencies, manage missing data more effectively, and provide a robust framework for accounting for unobserved confounders. Traditional survival models like the Cox model [10] and survival random forests [17] often assume simple covariate independence or require the model to implicitly learn complex dependencies. In practice, however, covariates often exhibit nonlinear and high-order dependencies. Neglecting these dependencies can bias treatment effect estimates and degrade predictive performance. Although recent HTE methods [9,12] offer greater flexibility through tree-based techniques, they typically rely on the assumption of covariate independence in their partitioning schemes. In survival analysis,  [17] develop random survival forests for right censored data, while [12] extend causal forests for HTE estimation. Alternatively,  [18] propose DeepHit, a neural network that directly models survival time distributions without assuming an underlying stochastic process. These models often use evaluation metrics, such as the concordance index [19], Brier scores, and calibration plots to assess treatment effect estimation performance.

To address these limitations, we propose incorporating copula-based transformations into survival models to explicitly model and exploit complex feature dependencies, thereby improving the estimation of HTE under-censoring. Our main contribution lies in the development of a Copula-Based Random Forest (CBRF) framework for survival analysis, where copulas model feature dependencies before constructing the random forest [16]. By explicitly addressing these dependencies, our approach improves both the interpretability and predictive performance of survival models in HTE settings. Investigating racial disparities in healthcare outcomes is especially important in survival contexts. While survival models can predict patient outcomes, their effectiveness often depends on how well they account for race, treatment status, and their interaction. Individuals from different racial groups may respond differently to treatment due to genetic, environmental, or socio-economic factors [20,21]. Incorporating such heterogeneity into survival models is thus essential for equitable predictive modeling.

In this study, we develop an integrated framework that combines deep learning models, specifically Long Short-Term Memory (LSTM) networks and hybrid Convolutional Neural Network (CNN)–LSTM models with a CBRF model. The LSTM model captures temporal dependencies in the data, while the CNN-LSTM hybrid extracts complex spatial and temporal features simultaneously. Meanwhile, the CBRF model captures nonlinear interactions, especially those involving race and treatment. We apply our methodology to breast cancer data from the TCGA-BRCA dataset, focusing on clinical features and gene expression data from Black or African American and White patients. We reduce the dimensionality of gene expression data using Principal Component Analysis (PCA) before model training. We evaluate all models using performance metrics including Root Mean Squared Error (RMSE), Mean Absolute Error (MAE), and the concordance index, along with treatment effect measures such as Average Treatment Effect (ATE) and Conditional Average Treatment Effect (CATE).

This research underscores the importance of incorporating HTE in survival models and demonstrates how race influences treatment response. Our findings suggest that machine learning models can support more equitable and personalized treatment recommendations, particularly in cancer treatment and survivorship planning [22].

The rest of the paper is structured as follows: Section 2 describes the dataset and the data collection process. Section 3 outlines the methodology, including preprocessing steps, model development, HT weight, IPTW, and evaluation criteria. Section 4 presents results, model performance metrics, and treatment effect estimates. Section 5 discusses the implications of our findings and suggests future research directions.

2. Differential Gene Expression Analysis of TCGA-BRCA Data

This study employs the DESeq2 R package to perform differential gene expression analysis on breast cancer data from the TCGA-BRCA dataset [23]. The analysis workflow begins with querying, downloading, and preparing gene expression data from the Genomic Data Commons (GDC), followed by constructing a dataset compatible with DESeq2 to identify Differentially Expressed Genes (DEGs). This workflow provides a systematic and reproducible framework for assessing gene expression changes in breast cancer samples, thereby offering valuable insights into the molecular mechanisms underlying tumor progression and treatment response.

Gene expression analyses of cancer datasets, such as those available through The Cancer Genome Atlas (TCGA), consistently yield significant insights into the molecular alterations associated with tumorigenesis [24,25,26]. In particular, RNA Seq-based differential gene expression analysis facilitates the identification of genes that are significantly upregulated or downregulated in cancerous tissues relative to normal tissues. DESeq2, a widely adopted tool for RNA Seq data analysis, normalizes raw counts and enables the identification of DEGs between conditions, such as tumor versus normal tissue [23].

To access the TCGA-BRCA dataset, we utilize the TCGAbiolinks package in R [24]. We construct a query targeting gene expression quantification data under the “Transcriptome Profiling” category for the TCGA-BRCA project. The raw RNA Seq count data are subsequently downloaded using the GDCdownload() function and structured using GDCprepare(), yielding a SummarizedExperiment object that integrates both the gene expression matrix and corresponding clinical metadata.

To maintain consistency and comparability in the downstream analysis, we filter the dataset based on the race metadata field, retaining only patients who self-identify as either Black or African American or White. We extract vital status information and encode it as a binary outcome (1 = deceased, 0 = alive). We remove genes with zero variance across samples to eliminate non-informative predictors.

Given the high dimensionality of transcriptomic data, we apply Principal Component Analysis (PCA) to reduce dimensionality while preserving the majority of variance. We transpose the gene expression matrix so that rows represent samples, and columns represent genes. Following variance filtering, we perform PCA using the prcomp() function in R, enabling centering and scaling. We retain the first six Principal Components (PCs), which capture the largest proportion of variance in gene expression, for use in subsequent predictive modeling.

We process and align clinical covariates such as race and vital_status with the gene expression data. After filtering, we subset the gene expression matrix to include only those samples for which clinical metadata are available. This alignment ensures that the rows of the reduced PCA matrix correspond exactly to the samples included in the clinical dataset. The final dataset comprises the following:

• A reduced gene expression matrix (samples × 6 PCs),
• A categorical variable indicating race (Black or African American vs. White),
• A binary outcome variable representing vital status.

We conduct all preprocessing steps in R using the following packages: TCGAbiolinks, DESeq2, dplyr, ggplot2, keras, randomForest, copula, Hmisc, and tidyr. These tools collectively support querying TCGA data, reducing dimensionality, handling clinical covariates, and preparing the dataset for advanced statistical modeling and machine learning applications.

Figure 1 presents the PCA screeplot for the breast cancer dataset. Based on this plot, we select six principal components for further analysis. The analysis begins by loading the gene expression and clinical data. The clinical dataset includes patient attributes such as race and vital status (alive or deceased). We categorize race into two groups: Black or African American and White, and filter the dataset to retain only these groups. We apply PCA to the gene expression data to reduce dimensionality and improve model training and interpretability.

Let $\mathbf{X}\in {\mathbb{R}}^{n\times m}$ denote the gene expression matrix, where n is the number of samples and m is the number of genes. After performing PCA, we project the data onto the first k principal components:

$${\mathbf{X}}_{\mathrm{reduced}}=\mathbf{X}·{\mathbf{V}}_k$$

Here, ${\mathbf{V}}_k$ contains the top k principal components, and ${\mathbf{X}}_{\mathrm{reduced}}\in {\mathbb{R}}^{n\times k}$ represents the reduced dataset, with · denoting matrix multiplication.

The preprocessing pipeline consists of the following steps:

• Convert the variable race into a factor to ensure appropriate handling in modeling.
• Encode the vital_status variable as a binary outcome (0 = alive, 1 = deceased) to facilitate survival analysis.
• Normalize and transform the gene expression matrix to a format compatible with statistical models.

This structured and filtered dataset serves as the foundation for downstream survival analysis and model development.

3. LSTM, CNN-LSTM, and Copula-Based Random Forest

Researchers widely adopt Long Short-Term Memory (LSTM) networks and CNN-LSTM models for time-series analysis, particularly in contexts that involve sequential dependencies. The Random Forest (RF) algorithm serves as an ensemble learning method based on decision trees. It trains multiple decision trees on bootstrapped samples and aggregates predictions through majority voting (for classification) or averaging (for regression). RF naturally handles nonlinearity, feature selection, and variable interactions. However, it assumes independence among input features, which often limits its effectiveness in scenarios with strong inter-feature dependencies.

Real-world datasets frequently contain complex, nonlinear dependencies that standard RF models fail to capture. To address this limitation, we enhance RF using copula transformations. This section introduces how we apply copula transformations to construct the Copula-Based Random Forest (CBRF) model, enabling structured dependency learning across features.

3.1. Gaussian Copula Transformation

The Gaussian copula transformation is a popular method for modeling multivariate dependence structures, especially when marginal distributions are non-normal. This approach decouples the modeling of marginal distributions from the joint dependence structure, making it a flexible tool for multivariate analysis. According to Sklar’s Theorem, any multivariate Cumulative Distribution Function (CDF) $F(x_1,x_2,\dots ,x_d)$ with continuous marginals $F_1,F_2,\dots ,F_d$ can be expressed in terms of a copula C as:

$$F(x_1,x_2,\dots ,x_d)=C(F_1\left(x_1\right),F_2\left(x_2\right),\dots ,F_d\left(x_d\right)),$$

(1)

where the copula $C:{[0,1]}^d\to [0,1]$ captures the dependence structure between variables, independent of their marginal distributions [27].

Given a d-dimensional random vector $\mathbf{X}=(X_1,X_2,\dots ,X_d)$, each marginal $X_j$ is transformed to the unit interval using its CDF:

$$U_j=F_j\left(X_j\right),j=1,\dots ,d.$$

(2)

In practice, when the true marginal distributions are unknown, empirical cumulative distribution functions ${\widehat{F}}_j$ are used, resulting in pseudo-observations ${\widehat{U}}_j={\widehat{F}}_j\left(X_j\right)$. These transformed variables $\mathbf{U}=(U_1,\dots ,U_d)$ are uniformly distributed on ${[0,1]}^d$, marginally.

To apply a Gaussian copula, the uniformly distributed variables are then transformed to standard normal marginals via the probit function, i.e., the inverse standard normal CDF:

$$Z_j={\Phi }^{-1}\left(U_j\right),j=1,\dots ,d,$$

(3)

where $\Phi$ denotes the standard normal cumulative distribution function. The resulting vector $\mathbf{Z}=(Z_1,\dots ,Z_d)$ has standard normal marginals and retains the dependence structure of the original data under the Gaussian copula assumption.

The dependence structure of the transformed data $\mathbf{Z}$ is characterized by estimating the empirical correlation matrix $\widehat{\Sigma }$. The Gaussian copula $C_{\Sigma }$ is defined as:

$$C_{\Sigma }(u_1,\dots ,u_d)={\Phi }_{\Sigma }\left({\Phi }^{-1}\left(u_1\right),\dots ,{\Phi }^{-1}\left(u_d\right)\right),$$

(4)

where ${\Phi }_{\Sigma }$ is the joint CDF of a multivariate normal distribution with zero mean and correlation matrix $\Sigma$. This copula captures the full joint dependence while preserving the marginal distributions of the original data.

The Gaussian copula transformation is useful in applications where complex marginal distributions need to be modeled while assuming a multivariate normal-like dependence structure. It has been widely applied in fields, such as financial econometrics, survival analysis, and risk management, and serves as a foundation for copula-based regression and deep learning models that capture non-linear and non-Gaussian dependencies [28].

In survival analysis, censored data (e.g., when subjects drop out or the study ends before the event of interest occurs) are common. Copulas are effective tools for modeling such data, as they handle both censored and uncensored data by capturing the dependence structure between them. This allows for more accurate estimation of treatment effects, even when some data are censored [29].

Moreover, in treatment effect estimation, copulas provide a powerful method for modeling non-linear, dependent, and higher-order relationships among variables, especially in HTE settings. By using copulas, we can obtain more reliable treatment effect estimates, particularly when other models struggle to account for complex dependencies [30].

3.2. Copula-Based Random Forest Model

Common copula families are Gaussian copula (captures symmetric dependencies), Clayton copula (models asymmetric lower tail dependence), Gumbel copula (captures upper tail dependence), Frank copula (captures both positive and negative dependence) and empirical copula (nonparametric, data-driven approach). We propose CBRF, a hybrid model that integrates copula functions with the traditional random forest algorithm to capture nonlinear dependencies and complex interactions between features in high-dimensional data. The key idea in this paper is to use the Gaussian copula function to model dependence structures before feeding the transformed features into a Random Forest (RF) model [31].

This approach enhances the flexibility and robustness of random forests, especially in scenarios where traditional independence assumptions are not valid. It is particularly useful for HTE estimation, financial risk modeling, survival analysis, and spatial data modeling.

The CBRF modifies standard RF by transforming the feature space using copulas before feeding it into the RF model. In this research, we apply the Gaussian copula to RF for data analysis. CBRF combines the concepts of copula theory and random forests to model complex dependencies between variables in multivariate data. It is particularly useful for dealing with high-dimensional datasets where the relationship between variables is intricate and nonlinear.

Random forests are ensembles of decision trees trained using random subsets of data and features. For regression tasks, each tree in the forest provides a prediction, and the final prediction is typically the average of all tree predictions:

$$f_{\mathrm{RF}}\left(X\right)={\displaystyle \frac{1}{T}}\sum _{t=1}^T{h}_t\left(X\right)$$

(5)

where $h_t\left(X\right)$ is the prediction from the t-th tree and T is the total number of trees in the forest [31].

The CBRF model integrates Gaussian copula transformations with random forests to effectively model complex, non-linear dependencies among features, particularly in high-dimensional and heterogeneous datasets. This hybrid approach benefits from the flexibility of the copula framework and the robustness of ensemble learning.

The Gaussian copula enables the modeling of multivariate dependencies through bivariate copulas while preserving the marginal distributions of the original variables. To implement the CBRF, pseudo-observations are first obtained from the empirical marginal distributions. These are then transformed into a standard normal scale using the inverse Gaussian CDF. The transformed features are used to train a random forest model.

The training process involves building decision trees using splitting criteria based on the transformed feature space. The final prediction is obtained by averaging the predictions across all trees:

$$\widehat{y}={\displaystyle \frac{1}{N_T}}\sum _{t=1}^{N_T}{T}_t\left(x\right)$$

(6)

where $T_t\left(x\right)$ is the prediction from the t-th tree for input x, and $N_T$ is the total number of trees.

This methodology captures both marginal behavior and inter-variable dependencies while maintaining the non-parametric, data-driven strengths of random forests. The randomForest() function is employed to train the model using bootstrapped samples of the copula-transformed data. The ensemble of trees captures non-linear patterns without requiring parametric assumptions, thus enhancing predictive performance.

3.3. LSTM and CNN-LSTM

LSTM networks use memory cells to retain long-term dependencies in time-series data. The main equations governing LSTM cell states are the following:

$$\begin{array}{cc}\hfill f_t& =\sigma (W_f[h_{t-1},x_t]+b_f),\hfill \\ \hfill i_t& =\sigma (W_i[h_{t-1},x_t]+b_i),\hfill \\ \hfill {\tilde{C}}_t& =tanh(W_C[h_{t-1},x_t]+b_C),\hfill \\ \hfill C_t& =f_t·C_{t-1}+i_t·{\tilde{C}}_t,\hfill \\ \hfill o_t& =\sigma (W_o[h_{t-1},x_t]+b_o),\hfill \\ \hfill h_t& =o_t·tanh\left(C_t\right).\hfill \end{array}$$

Here, $f_t$ is the forget gate, $i_t$ is the input gate, $C_t$ is the cell state, and $h_t$ is the hidden state.

The LSTM is a type of Recurrent Neural Network (RNN) designed to handle sequential data. LSTM cells are known for their ability to retain long-term dependencies, making them suitable for time-series forecasting, gene expression data, and other tasks involving sequential data. The LSTM model takes a 3D input of shape (samples, time_steps, features). In this case, the reshaped expression data (copula_transformed_reshaped) have dimensions $(n\_samples,1,n\_features)$, where $n\_samples$ is the number of data points (patients), 1 is the time step (since it is not a temporal sequence in the classical sense, but each sample is treated as a sequence of length 1), $n\_features$ is the number of principal components. The core of the model is the LSTM layer with 64 units. It processes the input sequence and learns patterns in the data. The return_sequences = FALSE argument indicates that only the last output (the hidden state at the final time step) is passed to the next layer. This is followed by a dense layer. After the LSTM layer, there is a dense layer with 32 units and ReLU activation. This layer introduces non-linearity and learns more complex relationships between the features. The final layer is a dense layer with a single neuron that outputs the predicted outcome (e.g., the vital status). The model is compiled with the MSE loss function, which is appropriate for regression tasks, where the output is continuous. The Adam optimizer is used to adjust the model weights during training. Adam is an adaptive optimizer that works well for a variety of tasks. The model is trained using the fit() method, where the copula-transformed data are used as input, and the outcome (vital status) is the target variable. Training is performed for 20 epochs, with a batch size of 32, and 20% of the data are reserved for validation.

The CNN-LSTM model is a hybrid deep learning architecture that combines Convolutional Neural Networks (CNNs) and LSTM networks. This architecture leverages the strengths of both CNNs (for spatial feature extraction) and LSTMs (for temporal sequence modeling), making it highly suitable for tasks involving spatiotemporal data such as time series forecasting, video classification, and functional data analysis. The CNN component is designed to extract high-level local features from the input sequence. For time series data, 1D convolutions are typically employed to identify short-term dependencies and localized temporal patterns across fixed-size windows. The LSTM network models long-term dependencies and sequential dynamics by utilizing memory cells regulated by input, forget, and output gates. This allows the model to retain important temporal information over extended time horizons. The CNN-LSTM architecture proceeds as follows: The raw sequential input $\mathbf{X}=\{x_1,x_2,\dots ,x_T\}\in {\mathbb{R}}^{T\times d}$ is segmented into fixed-size (possibly overlapping) windows. CNN Layers are processed through one or more 1D convolutional layers with kernel size k, resulting in feature maps that capture localized temporal structure. The resulting feature maps are reshaped as needed and passed through LSTM layers, which model the temporal dynamics of the CNN-extracted features. The LSTM output (either the last hidden state or the entire sequence, depending on the task) is fed into one or more fully connected (dense) layers to produce the final output.

Let $\mathrm{Conv}\left(\mathbf{X}\right)$ denote the CNN output and $\mathrm{LSTM}(·)$ the LSTM operation. The CNN-LSTM model can be expressed as follows:

$$\mathbf{H}=\mathrm{LSTM}\left(\mathrm{Conv}\right(\mathbf{X}\left)\right)$$

(7)

$$\widehat{\mathbf{y}}=\sigma (\mathbf{W}\mathbf{H}+\mathbf{b})$$

(8)

where $\mathbf{H}$ is the hidden representation from the LSTM, $\widehat{\mathbf{y}}$ is the predicted output, $\sigma$ is an activation function (e.g., softmax or identity), $\mathbf{W}$ and $\mathbf{b}$ are learnable weight and bias parameters.

CNN captures local patterns, while LSTM models global temporal dependencies. Convolutional layers reduce input dimensionality before temporal modeling for efficient handling of high-dimensional and long-sequence data. The CNN-LSTM model combines CNNs and LSTMs to capture both local patterns (through CNNs) and sequential dependencies (through LSTMs). Like the LSTM model, the CNN-LSTM model takes a 3D input tensor of shape (samples, time_steps, features) with dimensions $(n\_samples,1,n\_features)$. A 1D convolutional layer with 64 filters and a kernel size of 1 is applied to the input data to extract local patterns. The kernel size of 1 allows the model to learn spatial dependencies among features across samples. ReLU activation is used to introduce non-linearity. A max pooling layer with a pool size of 1 is applied, but it does not change the data size due to the pool size being equal to 1. The LSTM layer processes the sequential dependencies of the data with 64 units. Like the LSTM model, return_sequences = FALSE ensures that only the final hidden state is passed to the next layer. A dense layer with 32 units and ReLU activation is applied to further process the features learned by the CNN and LSTM layers. The final dense layer with a single neuron outputs the predicted value (vital status). The model is compiled with MSE as the loss function. The Adam optimizer is used to update the model weights. The model is trained similarly to the LSTM model using the fit() method, with 20 epochs and a batch size of 32, where 20% of the data is used for validation.

In the CNN-LSTM architecture, we employ a convolutional kernel of size 1 to perform pointwise transformations. This design choice allows the model to reweight features at each time step without aggregating information across adjacent steps, thereby preserving the temporal resolution of the input sequence. The pooling size is also set to 1, which maintains the sequence length across layers. Although such configurations do not introduce spatial abstraction, they are beneficial in scenarios where per-time-step feature transformation is prioritized over downsampling. Future work may consider alternative kernel and pooling sizes to explore their effects on feature summarization and model performance.

3.4. HT and IPTW Weights in Causal Inference

In causal inference, HT weights are used to correct for selection bias and estimate causal effects when treatment assignment is not randomized. These weights originate from survey sampling and are adapted for use under the potential outcomes framework.

The HT weight for each unit is the inverse of the probability of receiving the treatment actually received:

$$w_i=\left\{\begin{array}{cc}{\displaystyle \frac{1}{e\left(X_i\right)}}\hfill & \mathrm{if}{T}_i=1\hfill \\ {\displaystyle \frac{1}{1-e\left(X_i\right)}}\hfill & \mathrm{if}{T}_i=0\hfill \end{array}\right.$$

where $w_i$ is HT weight for unit i, $T_i$: treatment indicator (1 for treated, 0 for control), and $e\left(X_i\right)=P(T_i=1\mid X_i)$: propensity score, the probability of receiving the treatment given covariates $X_i$.

HT weights allow for reweighting the sample to approximate a randomized experiment, balancing the covariate distribution across treatment groups.

HT weights are commonly used to estimate the Average Treatment Effect (ATE):

$${\widehat{\mathrm{ATE}}}_{\mathrm{HT}}={\displaystyle \frac{1}{n}}\sum _{i=1}^n\left({\displaystyle \frac{T_i{Y}_i}{e\left(X_i\right)}}-{\displaystyle \frac{(1-T_i)Y_i}{1-e\left(X_i\right)}}\right)$$

This estimator is unbiased under the following assumptions: No unmeasured confounders and each unit has a nonzero probability of receiving either treatment.

HT weights can be unstable when propensity scores are close to 0 or 1. In practice, stabilized weights or trimming/extreme value handling is often used.

Inverse Probability of Treatment Weighting (IPTW) is a method for estimating causal effects from observational data by constructing a weighted pseudo-population in which treatment assignment is independent of measured covariates. This method relies on estimating the probability of the treatment conditional on observed covariates referred to as the propensity score and using its inverse as a weight.

Let $T_i\in \{0,1\}$ denote the treatment indicator for unit i, where $T_i=1$ indicates treatment and $T_i=0$ indicates control. Let $Y_i$ denote the observed outcome, and $X_i\in {\mathbb{R}}^p$ represent a vector of observed pre-treatment covariates. The propensity score is defined as:

$$e\left(X_i\right)=P(T_i=1\mid X_i),$$

i.e., the probability that unit i receives the treatment, given covariates $X_i$. Assuming that $e\left(X_i\right)$ is known or consistently estimated from the data, the IPTW weight for individual i is given by:

$$w_i^{\mathrm{IPTW}}={\displaystyle \frac{T_i}{e\left(X_i\right)}}+{\displaystyle \frac{1-T_i}{1-e\left(X_i\right)}}.$$

In practice, these weights correspond to:

$$w_i^{\mathrm{IPTW}}=\left\{\begin{array}{cc}{\displaystyle \frac{1}{e\left(X_i\right)}}\hfill & \mathrm{if}{T}_i=1,\hfill \\ {\displaystyle \frac{1}{1-e\left(X_i\right)}}\hfill & \mathrm{if}{T}_i=0.\hfill \end{array}\right.$$

These weights are used to create a weighted sample in which the distribution of observed covariates is balanced between treatment groups, thereby mimicking a randomized controlled trial.

Using these weights, the IPTW estimator of the Average Treatment Effect (ATE) is defined as:

$${\widehat{\mathrm{ATE}}}_{\mathrm{IPTW}}={\displaystyle \frac{1}{n}}\sum _{i=1}^n\left({\displaystyle \frac{T_i{Y}_i}{e\left(X_i\right)}}-{\displaystyle \frac{(1-T_i)Y_i}{1-e\left(X_i\right)}}\right).$$

This estimator is consistent for the true ATE under the following assumptions: $\left(Y\right(0),Y(1\left)\right)\perp T\mid X$, i.e., there are no unmeasured confounders given X, $0<e\left(X_i\right)<1$ for all i, ensuring that each unit has a positive probability of receiving either treatment. The model used to estimate the propensity scores $e\left(X_i\right)$ must be correctly specified.

In applications where extreme weights can introduce high variance, stabilized or truncated versions of IPTW are commonly used to enhance the robustness and efficiency of the estimator.

3.5. Model Training and Evaluation

All models are trained to predict the binary vital_status outcome using the same training dataset. Each model is compiled with the Adam optimizer and mean squared error loss function. Training includes a validation split to monitor generalization and prevent overfitting and runs for 20 epochs.

Model performance is assessed using Root Mean Squared Error (RMSE), Mean Absolute Error (MAE), and the Concordance statistic (C-index). RMSE and MAE evaluate prediction accuracy by measuring the average magnitude of errors:

$$\mathrm{RMSE}=\sqrt{{\displaystyle \frac{1}{n}}\sum _{i=1}^n{(y_i-{\widehat{y}}_i)}^2},\mathrm{MAE}={\displaystyle \frac{1}{n}}\sum _{i=1}^n|y_i-{\widehat{y}}_i|$$

To assess the model’s discriminatory power in survival outcomes, we use the C-statistic, also known as the Concordance index (C-index), is a metric used to evaluate the discriminatory power of a predictive model, particularly in survival analysis. It measures how well the model distinguishes between subjects with different survival outcomes based on their predicted risk scores. It provides an overall assessment of the model’s ability to correctly rank the subjects according to their likelihood of experiencing the event of interest (e.g., death, failure, etc.). The C-statistic evaluates the degree of concordance between predicted risk scores and the actual observed outcomes. It is calculated as the proportion of all possible pairs of subjects for which the predicted risk scores are in agreement with the actual observed event times. The C-statistic measures how well a model can rank subjects based on predicted risk scores relative to actual survival times. It evaluates the model’s ability to correctly order survival times.

For two individuals i and j with survival times $T_i,T_j$ and predicted risk scores ${\widehat{r}}_i,{\widehat{r}}_j$, the C-statistic is calculated as:

$$C={\displaystyle \frac{1}{\left|\mathcal{C}\right|}}\sum _{(i,j)\in \mathcal{C}}\left\{\begin{array}{cc}1,\hfill & \mathrm{if}{\widehat{r}}_i>{\widehat{r}}_j\mathrm{and}{T}_i<T_j\hfill \\ 0.5,\hfill & \mathrm{if}{\widehat{r}}_i={\widehat{r}}_j\mathrm{and}{T}_i\ne T_j\hfill \\ 0,\hfill & \mathrm{otherwise}\hfill \end{array}\right.$$

where $\mathcal{C}$ is the set of comparable pairs. A C-statistic of 0.5 suggests random predictions, while 1.0 indicates perfect concordance [19].

To capture treatment heterogeneity, we compute both the Average Treatment Effect (ATE) and the Conditional Average Treatment Effect (CATE). ATE is defined as the difference in expected outcomes between treated and control groups:

$$\mathrm{ATE}=\mathbb{E}\left[Y\right(1\left)\right]-\mathbb{E}\left[Y\right(0\left)\right]$$

In our context, race is treated as a moderator, with ATE computed separately for Black or African American and White populations.

CATE is the expected treatment effect conditional on covariates X, represented here by the first principal component score from a PCA decomposition:

$$\mathrm{CATE}\left(X\right)=\mathbb{E}\left[Y\right(1)-Y(0)\mid X]$$

This allows us to evaluate how treatment effects vary within racial groups based on individual-level covariate patterns.

4. Results

To begin our analysis, we examine the racial composition and survival outcomes within the dataset.

Table 1. Counts of race groups and survival outcomes.

Race	Alive	Dead	Total
Black or African American	159	32	191
White	721	159	880
Total	880	191	1071

summarizes the counts of patients who are alive or deceased across two racial groups: Black or African American and White.

Based on the data in Table 1, we calculate the marginal probabilities of race group membership, which are later used in the computation of Horvitz–Thompson (HT) weights:

$$P(T=1)={\displaystyle \frac{191}{1071}}\approx 0.1784,P(T=0)={\displaystyle \frac{880}{1071}}\approx 0.8216$$

Using these probabilities, the HT weights for each group are computed as:

$$w={\displaystyle \frac{1}{P\left(T\right)}}$$

• For Black or African American ($T=1$):

  $$w={\displaystyle \frac{1}{0.1784}}\approx 5.605$$
• For White ($T=0$):

  $$w={\displaystyle \frac{1}{0.8216}}\approx 1.217$$

These results are presented in

Table 2. Horvitz–Thompson weights.

Race	HT Weight
Black or African American	5.605
White	1.217

.

Next, we compute the conditional probabilities of treatment (race) given survival outcomes. These are necessary for estimating the Inverse Probability of Treatment Weighting (IPTW):

$$P(T=1|\mathrm{Alive})={\displaystyle \frac{159}{880}}\approx 0.1807,P(T=1|\mathrm{Dead})={\displaystyle \frac{32}{191}}\approx 0.1675$$

$$P(T=0|\mathrm{Alive})={\displaystyle \frac{721}{880}}\approx 0.8193,P(T=0|\mathrm{Dead})={\displaystyle \frac{159}{191}}\approx 0.8325$$

Accordingly, the IPTW values are calculated as follows:

$$IPTW={\displaystyle \frac{1}{P\left(T\right|Y)}}$$

The computed IPTW values are summarized in

Table 3. Inverse Probability of Treatment Weights (IPTW).

Race	Outcome	Probability $\mathit{P}\left(\mathit{T}\right|\mathit{Y})$	IPTW $1/\mathit{P}\left(\mathit{T}\right|\mathit{Y})$
Black or African American	Alive	0.1807	5.53
Black or African American	Dead	0.1675	5.97
White	Alive	0.8193	1.22
White	Dead	0.8325	1.20

.

The application of HT and IPTW weights plays a critical role in causal inference for observational studies. HT weights adjust for heterogeneity in treatment effects across subgroups, allowing for individualized treatment effect estimates. IPTW, in contrast, corrects for selection bias by reweighting individuals based on the likelihood of receiving their observed treatment, given covariates. This process emulates random treatment assignment and helps balance the covariate distribution between treatment groups.

Using both weights in tandem allows for a more accurate and equitable estimation of treatment effects—HT weights enhance precision by addressing within-group variation, while IPTW improves validity by mitigating confounding bias.

To evaluate model performance and the effectiveness of the weighting strategies, we compile results from CNN-LSTM, LSTM, and CBRF models into a comprehensive summary table. This table includes performance metrics, such as RMSE, MAE, and the C-statistic, as well as estimated ATE and race-specific CATE.

To aid interpretability, we visualize these outcomes using bar plots that disaggregate metrics by model and racial group. These visualizations clearly illustrate the comparative effectiveness of each model and highlight disparities in treatment effects across racial subgroups.

Table 4. Model performance comparison across weighting methods.

Model	Race	RMSE	MAE	C_Statistic	ATE	CATE	Weighting
CNN-LSTM	Black or African American	0.2991	0.2147	0.8890	1.1980	1.2080	No Weights
CNN-LSTM	White	0.2991	0.2147	0.8890	1.2073	1.2011	No Weights
CBRF	Black or African American	0.5037	0.4913	0.5042	1.4638	1.4883	No Weights
CBRF	White	0.5037	0.4913	0.5042	1.4924	1.4891	No Weights
LSTM	Black or African American	0.3112	0.2330	0.8853	1.2204	1.2142	No Weights
LSTM	White	0.3112	0.2330	0.8853	1.2109	1.2014	No Weights
CNN-LSTM	Black or African American	0.2916	0.2047	0.8893	1.2035	1.1624	HW Weights
CNN-LSTM	White	0.2916	0.2047	0.8893	1.1889	1.2116	HW Weights
CBRF	Black or African American	0.4829	0.4714	0.5140	1.4366	1.4527	HW Weights
CBRF	White	0.4829	0.4714	0.5140	1.4813	1.5294	HW Weights
LSTM	Black or African American	0.3041	0.2036	0.8706	1.1325	1.1569	HW Weights
LSTM	White	0.3041	0.2036	0.8706	1.1676	1.2120	HW Weights
CNN-LSTM	Black or African American	0.2604	0.1745	0.8944	1.0971	1.1325	IPTW Weights
CNN-LSTM	White	0.2604	0.1745	0.8944	1.1485	1.1079	IPTW Weights
CBRF	Black or African American	0.4972	0.4889	0.4885	1.4748	1.4816	IPTW Weights
CBRF	White	0.4972	0.4889	0.4885	1.4978	1.4855	IPTW Weights
LSTM	Black or African American	0.2751	0.1894	0.8748	1.1214	1.1558	IPTW Weights
LSTM	White	0.2751	0.1894	0.8748	1.1802	1.1479	IPTW Weights

presents a detailed comparison of model performance and treatment effect estimates across three machine learning models—CNN-LSTM, LSTM, and CBRF—evaluated under three weighting strategies: no weighting (unadjusted), HT weights, and IPTW. The performance metrics include RMSE, MAE, and C-statistic, which assess prediction accuracy and discrimination. Additionally, ATE and CATE are reported for two racial subgroups: Black or African American and White.

CNN-LSTM demonstrates superior predictive capability across all weighting methods. For the unweighted model, CNN-LSTM achieves an RMSE of 0.2991, MAE of 0.2147, and a C-statistic of 0.8890 across both racial groups. These values are consistently lower than those of LSTM (RMSE = 0.3112, MAE = 0.2330, C-statistic = 0.8853) and substantially better than CBRF (RMSE = 0.5037, MAE = 0.4913, C-statistic = 0.5042). When incorporating HT weights, CNN-LSTM shows marginal improvement in both error metrics (RMSE = 0.2916, MAE = 0.2047), while maintaining a stable C-statistic (0.8893). The most notable gains appear under IPTW weighting, where CNN-LSTM achieves its best overall performance (RMSE = 0.2604, MAE = 0.1745, C-statistic = 0.8944). LSTM also benefits from IPTW, improving to RMSE = 0.2751 and MAE = 0.1894, though still trailing CNN-LSTM.

CBRF, by contrast, maintains relatively poor performance across all weighting schemes. Despite slight improvements under HT (RMSE = 0.4829) and IPTW (RMSE = 0.4972), its C-statistic remains close to 0.50, suggesting limited discriminatory ability, and indicating performance near chance level.

ATE and CATE estimates vary both across models and racial subgroups. Without weighting, CNN-LSTM estimates ATEs of 1.1980 (Black or African American) and 1.2073 (White), with CATEs closely aligned. LSTM produces slightly higher ATEs (1.2204 for Black or African American; 1.2109 for White), suggesting consistent treatment effects across groups.

Under HT weighting, estimates shift moderately. CNN-LSTM reports an ATE of 1.2035 for Black or African American and 1.1889 for White, while CATEs exhibit greater variability (1.1624 for Black or African American; 1.2116 for White), indicating some heterogeneity. IPTW weighting further sharpens subgroup differences. CNN-LSTM yields an ATE of 1.0971 for Black or African American and 1.1485 for White, with CATEs of 1.1325 and 1.1079, respectively. These results suggest the potential presence of race-specific treatment response heterogeneity, with a marginally higher CATE for Black or African American individuals under the best-performing model configuration.

CBRF consistently produces higher ATE and CATE estimates across all settings. For instance, under IPTW, CBRF estimates an ATE of 1.4978 and a CATE of 1.4855 for White individuals—values notably higher than those produced by the deep learning models. However, given the model’s poor predictive performance, these estimates should be interpreted with caution and may reflect overfitting or bias rather than true effect magnitudes.

The choice of weighting strategy has a measurable impact on both model performance and effect estimation. Unweighted models yield higher variance and poorer calibration. HT weights provide moderate improvements, particularly in LSTM and CNN-LSTM, suggesting partial adjustment for covariate imbalance. IPTW consistently delivers the best results in terms of predictive accuracy and discrimination, particularly for CNN-LSTM. This supports the theoretical advantages of IPTW in mitigating confounding by balancing covariates across treatment groups.

These findings demonstrate that CNN-LSTM, particularly when combined with IPTW, offers the most reliable framework for treatment effect estimation in racially stratified cohorts. Its low error rates and high C-statistics across subgroups suggest both strong predictive validity and equitable performance. CBRF, while flexible in handling complex interactions, shows substantial limitations in generalizability due to poor prediction accuracy and model calibration.

Overall, Table 4 highlights the critical role of model selection and appropriate weighting in achieving accurate, robust, and equitable treatment effect estimates. A visualization of these results is provided in Figure 2, facilitating intuitive comparison across models, weighting strategies, and racial subgroups. The related R code can be found in Appendix A.

5. Conclusions

This study evaluated the comparative performance of advanced machine learning models—CNN-LSTM, LSTM, and CBRF—in estimating treatment effects under varying weighting strategies in a racially heterogeneous population. The results underscore the superior predictive and discriminatory performance of the CNN-LSTM model, particularly when integrated with IPTW. This model consistently demonstrated the lowest prediction errors and the highest C-statistics across racial subgroups, indicating its suitability for both outcome prediction and causal inference.

The use of IPTW substantially enhanced model calibration and reduced error variance, reinforcing its utility in addressing confounding in observational data. While HT weights offered modest gains over unweighted models, IPTW delivered the most reliable and equitable estimates of both ATE and CATE. These improvements were especially evident in the CNN-LSTM and LSTM models, highlighting the value of integrating modern weighting strategies with deep learning architectures.

Although CBRF estimated consistently higher treatment effects, its inferior predictive performance and poor discrimination capacity raise concerns about the validity and stability of these estimates. This suggests that, in contexts requiring both prediction and causal inference, model interpretability should not come at the expense of statistical rigor and generalization capability.

Importantly, the results revealed subtle yet meaningful heterogeneity in treatment effects across racial groups, which were most reliably captured by the CNN-LSTM model under IPTW. These findings advocate for the use of deep learning frameworks—especially CNN-LSTM paired with IPTW—in producing accurate and fair treatment effect estimates in racially diverse settings.

In sum, this study contributes to the growing literature on machine learning for causal inference by demonstrating the importance of model selection, weighting strategy, and subgroup-specific evaluation in health and policy research. Future work should extend these findings by validating performance on additional datasets, incorporating longitudinal and clustered data structures, and evaluating fairness metrics more explicitly.