Robust propensity score estimation via loss function calibration

Abstract

Propensity score estimation is often used as a preliminary step to estimate the average treatment effect with observational data. Nevertheless, misspecification of propensity score models undermines the validity of effect estimates in subsequent analyses. Prediction-based machine learning algorithms are increasingly used to estimate propensity scores to allow for more complex relationships between covariates. However, these approaches may not necessarily achieve covariates balancing. We propose a calibration-based method to better incorporate covariate balance properties in a general modeling framework. Specifically, we calibrate the loss function by adding a covariate imbalance penalty to standard parametric (e.g. logistic regressions) or machine learning models (e.g. neural networks). Our approach may mitigate the impact of model misspecification by explicitly taking into account the covariate balance in the propensity score estimation process. The empirical results show that the proposed method is robust to propensity score model misspecification. The integration of loss function calibration improves the balance of covariates and reduces the root-mean-square error of causal effect estimates. When the propensity score model is misspecified, the neural-network-based model yields the best estimator with less bias and smaller variance as compared to other methods considered.

1. Introduction:

In observational studies, the covariate distributions between treatment groups are usually imbalanced, which induces confounding bias when estimating the causal effect.^1,2 Propensity score (PS), defined as the conditional probability of being treated given the covariates, ³ plays an essential role in causal inference. Because of the balancing property, the covariates and treatment are conditionally independent given the PS. ⁴ There are several PS methods to remove the confounding bias when estimating causal effect, including weighting, ⁵ stratification, ⁶ matching, ⁷ or regression adjustment. ⁸

The standard procedure for propensity score estimation is to fit a logistic regression model, check covariate balance, and re-fit the model until the balance is achieved. ⁶ However, there are significant drawbacks to this approach. Firstly, the parametric form of a logistic regression could be easily misspecified.^9,10 It has been found that even under mild model misspecification, the logistic regression methods for PS estimation may fail and the performance of weighting methods including doubly robust estimator may be unsatisfactory. ¹¹ Secondly, the logistic regression may fail to capture the complex non-linear relationship between covariates and the treatment. ¹² Moreover, the iterative process is cumbersome, and people often skip this step, which leads to unreliable conclusions. ¹³

The goal of propensity score estimation is to balance covariates. Naïve PS estimation method such as a logistic regression may not achieve covariate balancing. ¹⁴ Empirical studies showed that the absolute standardized mean difference in covariates was strongly correlated with the bias of causal effect estimation, ¹⁵ but the prediction accuracy measure such as $c-$ statistic was not associated with the magnitude of bias. ¹⁶ Best prediction model may not achieve the covariate balance when there is model misspecification. Various approaches focusing on covariate balancing have been proposed, ¹⁷ including entropy balancing (EB), ¹⁸ stabilized balancing weights (SBWs), ¹⁹ covariate balancing propensity score (CBPS), ²⁰ covariate balancing scoring rules (CBSRs), ²¹ and propensity score local balance (PSLB). ²² Both EB and SBW methods directly model the weights required to balance covariates and bypass the propensity scores. Therefore, they are not applicable to the PS based methods, such as matching or stratification, for causal inference. The CBPS method estimates propensity scores using the logistic regression model, where the balance of covariates is optimized within the framework of generalized method-of-moments. ¹⁰ The CBSR method generalizes the concept of CBPS in various aspects, such as incorporating different PS link functions and applying non-parametric kernel methods. It can also handle different causal estimands while maintaining the covariate moment balancing constraints. ²¹ PSLB, an extension of CBPS method, considers the non-parametric kernel model and optimizes the local balance within the propensity score stratified (PS-stratified) sub-population. ²²

Statistical or machine learning approaches are more flexible to incorporate non-linear relationships. ²³ Prediction-based machine learning algorithms, such as bagging or boosting, ²⁴ random forest, boosted classification and regression trees, ¹² and neural networks are increasingly used to predict the propensity scores.^15,16 However, these approaches may not necessarily achieve covariates balancing. ²⁵ Incorporating covariate balance into the optimization process of machine learning methods can improve the estimation of causal effects.

Inspired by the dual characteristics of the PS which is the covariate balancing score and the conditional probability of treatment assignment, we calibrate the likelihood or loss function of a prediction model by adding a covariate-imbalance penalty. The calibrated loss function can take into account the prediction accuracy and covariate balance at the same time. This approach may mitigate the impact of misspecification of the PS model and thus improve the performance of the ATE estimator. Furthermore, we leverage the flexibility of machine learning methods by configuring the base model as either a logistic model or a neural network architecture.

The remainder of this article is organized as follows. In Section 2, we begin with a brief review of the traditional logistic regression PS model. Then we present the proposed methods where the loss function is calibrated by adding an imbalance score. Optimization and tuning parameter selection methods are introduced with a comprehensive algorithm. In Section 3, we assess the performance of the proposed methods in comparison with other PS estimation methods in the simulation study. In Section 4, we apply the proposed methods to a real study that examined the causal effect of the utilization of right heart catheterization in the intensive care unit on the hospital length of stay. We provide conclusions and discussions in Section 5.

2. Methods

2.1. Notation and model specification

Suppose there is a random sample of $n$ individuals from a population, let $Z_i=0\text{ or }1$ denote the binary treatment assignment, ${\mathbf{X}}_{\mathbf{i}}$ be the $p$ -dimensional covariate vector for the $i$ th individual, where $x_{ik}$ denotes the $k$ th covariate of the $i$ th individual for $i=1,\dots ,n$ and $k=1,\dots ,p$ . According to Rosenbaum and Rubin, ³ we define the potential outcomes for each individual as $Y_i(t)$ under the treatment $t=0,1$ . The observed outcome is $Y_i=Z_i{Y}_i(1)+(1-Z_i)Y_i(0)$ . We observe the $n\times p$ covariate matrix $\mathbf{X}$ , the treatment vector $\mathbf{Z}$ , and the observed outcome vector $\mathbf{Y}$ for the sample. Our proposed work focused on the estimation of the average treatment effect (ATE): ${\tau }^{ATE}=\mathbb{E}[\mathbf{Y}(1)-\mathbf{Y}(0)]$ . The same framework could be easily extended for other causal estimands, such as ATE among treated (ATT). To ensure the ATE is identifiable, we made the classic assumptions as discussed by Angrist et al., ²⁶ including the stable unit treatment value assumption (SUTVA), no unmeasured confounders, and the positivity assumption. ²⁷

The PS of receiving treatment of interest can be defined as $e({\mathbf{X}}_{\mathbf{i}})=\mathbb{P}(Z_i=1\mid {\mathbf{X}}_{\mathbf{i}})$ . The estimation of PS can be conducted in many ways, including the frequently employed logistic regression or increasingly prevalent machine learning methods.^10,28,29 We will start the illustration of our proposed method with logistic regression as the base model for simplicity, and extend it to the case when the neural network method is used as a base model.

The logistic PS model can take the form of:

$$\log \frac{e({\mathbf{X}}_{\mathbf{i}})}{1-e({\mathbf{X}}_{\mathbf{i}})}={\mathbf{X}}_{\mathbf{i}}^{\mathbf{T}}\mathit{\beta }$$ (1)

where $\mathit{\beta }$ is a vector of the logistic regression coefficients. Then the inverse propensity weight for the $i$ th individual can be defined as a function of $\mathit{\beta }$ :

$$w(Z_i,{\mathbf{X}}_{\mathbf{i}};\mathit{\beta })=w(Z_i,e({\mathbf{X}}_{\mathbf{i}};\mathit{\beta }))=\frac{Z_i}{e({\mathbf{X}}_{\mathbf{i}};\mathit{\beta })}+\frac{1-Z_i}{1-e({\mathbf{X}}_{\mathbf{i}};\mathit{\beta })}$$ (2)

To achieve the goal of balancing the covariates while accurately predicting the treatment assignment, we propose the calibrated loss function for the PS model that accounts for the imbalance of covariates between treatment groups. Specifically, we consider the sum of the cross entropy, which is equivalent to the negative log-likelihood for the logistic regression base model, and the penalty term, namely imbalance score (IS), as defined below:

$$L(\mathit{\beta })=-\sum _{i=1}^n[Z_i\log e({\mathbf{X}}_{\mathbf{i}};\mathit{\beta })+(1-Z_i)\log (1-e({\mathbf{X}}_{\mathbf{i}};\mathit{\beta }))]+\lambda \text{ IS}(\mathbf{X},\mathbf{Z};\mathit{\beta })$$ (3)

where the overall IS is defined as the sum of the IS for each covariate:

$$\text{ IS}(\mathbf{X},\mathbf{Z};\mathit{\beta })=\sum _{k=1}^p{\text{ IS}}_k(\mathbf{X},\mathbf{Z};\mathit{\beta })$$ (4)

and $\lambda$ is the tuning parameter.

There are several ways to measure the difference between two covariate distributions, such as differences in moments (mean, variance, etc.), Kullback–Leibler (KL) divergence, and Kolmogorov–Smirnov (KS) test statistics. ³⁰ Here we describe one commonly used measure for an IS, the absolute standardized difference in mean^10,31,32:

$${\text{ IS}}_k(\mathbf{X},\mathbf{Z};\mathit{\beta })=\frac{\mid \frac{\sum _{i=1}^{n}w(Z_i,{\mathbf{X}}_{\mathbf{i}};\mathit{\beta })x_{ik}{Z}_i}{\sum _{i=1}^{n}w(Z_i,{\mathbf{X}}_{\mathbf{i}};\mathit{\beta })Z_i}-\frac{\sum _{i=1}^{n}w(Z_i,{\mathbf{X}}_{\mathbf{i}};\mathit{\beta })x_{ik}(1-Z_i)}{\sum _{i=1}^{n}w(Z_i,{\mathbf{X}}_{\mathbf{i}};\mathit{\beta })(1-Z_i)}\mid }{\sqrt{\frac{s_{1k}^2}{\sum _{i=1}^{n}w(Z_i,{\mathbf{X}}_{\mathbf{i}};\mathit{\beta })Z_i}+\frac{s_{0k}^2}{\sum _{i=1}^{n}w(Z_i,{\mathbf{X}}_{\mathbf{i}};\mathit{\beta })(1-Z_i)}}}$$ (5)

where $s_{Zk}$ is the weighted sample variance of the $k$ th covariate for group $Z=0,1$ .

In fact, the logistic regression can be viewed as the most basic form of a neural network, comprising solely one input and output layer with a sigmoid function as the activation function, denoted by $S(t)=\frac{\exp (t)}{1+\exp (t)}$ . The regression coefficient vector $\mathit{\beta }$ defines the weights of the covariates. The PS is then estimated by applying the weighted sum of the covariates to the sigmoid function, that is, $S\left(\sum _k{\mathbf{X}}_{\mathbf{i}\mathbf{k}}{\beta }_k\right)=S({\mathbf{X}}_{\mathbf{i}};\mathit{\beta })$ . Figure 1 shows the structure of a logistic regression model and Figure 2 showcases one simple structure of a neural network with one hidden layer. To estimate the parameters in the neural network base model using the proposed method, we substitute the conventional prediction-based loss function with our novel calibrated loss function. By minimizing this loss function, we can estimate all parameters involved in the neural network model. ³³ With these estimated parameters, the propensity scores can be estimated. A more complicated structure of a neural network can have multiple hidden layers with different structures, designed to model non-linear relationships and capture intricate patterns and relationships in the data.^34,35 The structure of the neural network and the number of neurons in each layer should be tailored to the complexity of the dataset. Several methods have been proposed for designing prediction-based neural networks.^36–38 These techniques can be used to initially determine the structure of neural networks. Once the structure is established, the calibrated loss function in equation (3) can be easily incorporated to optimize the parameters.

Figure 1.

Neural network architecture for logistic regression.

Figure 2.

Neural network architecture with a hidden layer.

2.2. Optimization

To minimize the calibrated loss function, we use Algorithm 1 described below to estimate the parameter $\mathit{\beta }$ . Specifically, we utilize the gradient descent algorithm to minimize the loss function and update the parameter estimates iteratively based on the following equation:

$${\mathit{\beta }}^{\mathbf{n}\mathbf{e}\mathbf{w}}\leftarrow \mathit{\beta }-\eta \mathrm{\nabla }L(\mathit{\beta })$$

where $\eta$ is the learning rate to control the step size of each update and $\mathrm{\nabla }$ is the gradient of the loss function. We update the estimate of parameter $\hat{\mathit{\beta }}$ until the algorithm converges or when the maximum iteration ( $I$ ) is reached.

The initial values of the parameter $\mathit{\beta }$ are selected through a prediction-based PS model, ensuring a reasonable starting point for optimization. While the algorithm may converge to local minima, it is close to the prediction-based estimation. The primary objective of our optimization is not necessarily to find the global minimum but rather to enhance performance within an acceptable range. There are some modified gradient descent algorithms that help find a global minimum even for non-convex problems including adaptive learning rate algorithms such as Adam ³⁹ and Adagrad, ⁴⁰ which adjust step sizes dynamically during the training stage, and stochastic gradient descent (SGD), ⁴¹ which introduces randomness to escape local minima.

2.3. Tuning parameter

Since the objective of PS estimation is not to achieve perfect model fitting, the traditional tuning parameter selection methods that emphasize prediction accuracy cannot be applied. Several researchers suggested that using the covariate balance measures such as the absolute standardized difference as tuning parameter selection criteria worked better in ATE estimation.^15,32,42,43 Motivated by this idea, we use grid search and select the tuning parameter $\lambda$ that gives the smallest IS as defined in equation (5).

Cross-validation (CV) is widely recognized as an effective technique for parameter tuning. ⁴⁴ In our context, the application of CV is also feasible, contingent upon available computational resources. It is important to note that if the $\lambda$ is zero, the proposed method will degenerate to its base model. Conversely, with a sufficiently large $\lambda$ , the method achieves a trade-off between prediction accuracy and covariate balance. Hence, we recommend exploring a broad range of values for $\lambda$ during the grid search process.

3. Simulation study

The simulation design introduced by Kang and Schafer ¹¹ has become a widely adopted benchmark setting for evaluating the performance of PS methods.^19,20 To test the robustness of our proposed method under model misspecification, we replicated this simulation setting and examined the empirical performance of the proposed method under different model misspecifications. In detail, we applied the calibrated loss function methods for PS estimation based on two base models, logistic regression and one-hidden-layer neural network in which the activation function is ReLU ⁴⁵ for the hidden layer and the sigmoid function for the output layer.

3.1. Estimators of causal effect

In the following, we briefly reviewed three ATE estimators used in our simulation study. The first two were inverse probability weighting estimators that were proven to be an unbiased estimator when the PS model was correctly specified. ³¹

• (1)
  Horvitz–Thompson (HT) estimator ⁴⁶ :

  $${\tau }_{\mathrm{HT}}=\frac{1}{n}\sum _{i=1}^n(\frac{Z_i{Y}_i}{e({\mathbf{X}}_{\mathbf{i}})}-\frac{(1-Z_i)Y_i}{1-e({\mathbf{X}}_{\mathbf{i}})})$$ (6)
• (2)
  Hájek estimator (also called stabilized weights) ⁴⁷ :

  $${\tau }_{\text{H}\acute{\mathrm{a}}\text{jek}}=\sum _{i=1}^n\frac{Z_i{Y}_i}{e({\mathbf{X}}_{\mathbf{i}})}/\sum _{i=1}^n\frac{Z_i}{e({\mathbf{X}}_{\mathbf{i}})}-\sum _{i=1}^n\frac{(1-Z_i)Y_i}{1-e({\mathbf{X}}_{\mathbf{i}})}/\sum _{i=1}^n\frac{1-Z_i}{1-e({\mathbf{X}}_{\mathbf{i}})}$$ (7)

  The Hájek estimator could be more efficient (with a smaller variance) than the HT estimator. ⁴⁷
• (3)
  Doubly robust (DR) estimator ⁴⁸ :

  $${\tau }_{\mathrm{DR}}=n^{-1}\sum _{i=1}^n[\frac{Z_i{Y}_i}{e({\mathbf{X}}_{\mathbf{i}})}-\frac{Z_i-e({\mathbf{X}}_{\mathbf{i}})}{e({\mathbf{X}}_{\mathbf{i}})}{m}_1({\mathbf{X}}_{\mathbf{i}})]-n^{-1}\sum _{i=1}^n[\frac{(1-Z_i)Y_i}{1-e({\mathbf{X}}_{\mathbf{i}})}+\frac{Z_i-e({\mathbf{X}}_{\mathbf{i}})}{1-e({\mathbf{X}}_{\mathbf{i}})}{m}_0({\mathbf{X}}_{\mathbf{i}})]$$ (8)

  where $m_Z({\mathbf{X}}_{\mathbf{i}})=\mathbb{E}(Y_i(Z)\mid {\mathbf{X}}_{\mathbf{i}})\text{ for }Z=0,1$ .

The DR estimator combines a PS model with an outcome model to yield a consistent estimator when one of the models is correctly specified.^49–52

3.2. Simulation set-up

We first generated four latent pre-treatment covariates ${\mathbf{X}}_i=(X_{i1},X_{i2},X_{i3},X_{i4})$ that followed independent standard normal distributions. The potential outcomes were then generated through two linear regressions:

$$Y_i(1)=210+27.4X_{i1}+13.7X_{i2}+13.7X_{i3}+13.7X_{i4}+{ϵ}_i$$ (9)

and

$$Y_i(0)=200+27.4X_{i1}+13.7X_{i2}+13.7X_{i3}+13.7X_{i4}+{ϵ}_i$$ (10)

where ${ϵ}_i\sim N(0,1)$ . Under this setting, the mean potential outcome if everyone was treated was 210, the mean potential outcome if everyone was not treated was 200 and the true ATE was 10.

The true PS was generated with a logistic regression model:

$$\log \frac{e({\mathbf{X}}_{\mathbf{i}})}{1-e({\mathbf{X}}_{\mathbf{i}})}=-X_{i1}+0.5X_{i2}-0.25X_{i3}-0.1X_{i4}$$ (11)

The treatment assignment was then generated from $Z_i\sim \text{Bernoulli}(e({\mathbf{X}}_{\mathbf{i}}))$ . The observed outcome was $Y_i=Z_i{Y}_i(1)+(1-Z_i)Y_i(0)$ . To introduce model misspecification, we transformed the ${\mathbf{X}}_{\mathbf{i}}$ to ${\mathbf{X}}_{\mathbf{i}}^{\mathbf{\ast }}=(X_{i1}^{*},X_{i2}^{*},X_{i3}^{*},X_{i4}^{*})$ as follows:

$$\begin{array}{rl}{X}_{i1}^{*}& =\exp (X_{i1}/2)\end{array}$$ (12)

$$\begin{array}{rl}{X}_{i2}^{*}& =X_{i2}/(1+\exp (X_{i1}))+10\end{array}$$ (13)

$$\begin{array}{rl}{X}_{i3}^{*}& =(X_{i1}{X}_{i3}/25+0.6{)}^3\end{array}$$ (14)

$$\begin{array}{rl}{X}_{i4}^{*}& =(X_{i2}+X_{i4}+20{)}^2\end{array}$$ (15)

We investigated the performance of our proposed methods for three ATE estimators as described above under different scenarios. Since only the PS model was required for the HT and Hájek estimators, we assessed the PS estimation methods under two scenarios: (1) PS model was misspecified using ${\mathbf{X}}_{\mathbf{i}}^{\mathbf{\ast }}$ and (2) PS model was correctly specified using ${\mathbf{X}}_{\mathbf{i}}$ . For the DR estimator, we also considered whether the outcome model was correctly specified under an incorrect or a correct PS model. Specifically, for the misspecified outcome model we applied linear regression with observed covariates ${\mathbf{X}}_{\mathbf{i}}^{\mathbf{\ast }}$ . In contrast, we used the latent covariates ${\mathbf{X}}_{\mathbf{i}}$ for the correctly specified outcome model.

For each scenario, we conducted $1000$ Monte Carlo simulations with a sample size $N=200$ and $500$ . We compared the bias and root-mean-square error (RMSE) of estimated ATE in equations (6) to (8) based on the following PS estimation methods:

GLM: Standard logistic regression

GLM+IS: Proposed calibration method with logistic regression as the base model

NNet: Neural network

NNet+IS: Proposed calibration method with neural network as the base model

CBPS: CBPS algorithm that extends logistic regression to achieve covariate balance

The structure of the neural network in NNet and NNet+IS was chosen to be one hidden layer of four neurons as described in Figure 2, considering that we have four covariates, and the data is not very complex, suggesting that a single hidden layer would be sufficient. We assessed and ensured convergence of the proposed methods by ensuring the stability of the loss function over the final $100$ iterations.

Furthermore, we calculated the standardized mean difference in the weighted observed covariates ${\mathbf{X}}_{\mathbf{1}}^{\mathbf{\ast }}$ to ${\mathbf{X}}_{\mathbf{4}}^{\mathbf{\ast }}$ , and latent covariates ${\mathbf{X}}_{\mathbf{1}}$ to ${\mathbf{X}}_{\mathbf{4}}$ , to assess the covariate balance between the two groups. The inverse propensity weights were estimated using the methods described above.

3.3. Simulation results

Table 1 presents the bias and RMSE of ATE estimation for different effect estimators under various PS estimation methods. We also reported the Monte Carlo standard error of bias and RMSE in Table 2. ⁵³ Figure 3 shows the boxplots of ATE estimates across 1000 simulation runs, truncated at the 0.1th and 99.9th percentiles.

Table 1.

Performance of three average treatment effect estimators based on different propensity score estimation methods.

		GLM		GLM+IS		NNet		NNet+IS		CBPS
		Bias	RMSE	Bias	RMSE	Bias	RMSE	Bias	RMSE	Bias	RMSE
n = 200	PS model is misspecified
	HT	22.469	149.114	$-$ 1.652	13.000	1.076	25.503	$-$ 0.895	12.466	$-$ 4.986	5.940
	Hájek	$-$ 2.038	8.827	$-$ 5.057	5.949	$-$ 3.747	6.687	$-$ 2.987	3.812	$-$ 5.062	5.963
	DR (outcome model misspecified)	$-$ 7.780	31.260	$-$ 5.008	5.924	$-$ 3.769	4.692	$-$ 2.925	3.788	$-$ 5.028	5.951
	DR (outcome model correctly specified)	0.006	0.386	0.010	0.178	0.009	0.206	0.010	0.180	0.010	0.177
	PS model is correctly specified
	HT	0.029	15.735	0.109	14.844	0.284	13.342	1.187	12.492	0.100	0.239
	Hájek	$-$ 0.431	3.281	0.002	0.178	$-$ 0.936	3.778	0.001	0.191	0.008	0.177
	DR (outcome model misspecified)	$-$ 0.020	2.345	0.242	1.822	$-$ 0.656	2.028	0.014	1.358	0.218	1.830
	DR (outcome model correctly specified)	0.007	0.179	0.007	0.177	0.009	0.186	0.010	0.184	0.008	0.177
n = 500	PS model is misspecified
	HT	39.977	350.917	$-$ 0.400	5.613	3.029	62.053	$-$ 1.580	5.379	$-$ 5.568	5.994
	Hájek	0.230	10.551	$-$ 5.574	5.988	$-$ 1.393	5.526	$-$ 2.018	2.485	$-$ 5.583	5.999
	DR (outcome model misspecified)	$-$ 12.156	93.065	$-$ 5.570	5.992	$-$ 2.446	4.284	$-$ 1.995	2.470	$-$ 5.573	5.992
	DR (outcome model correctly specified)	$-$ 0.021	0.817	$-$ 0.002	0.110	0.016	0.588	$-$ 0.001	0.109	$-$ 0.001	0.110
	PS model is correctly specified
	HT	$-$ 0.228	9.255	$-$ 0.149	8.172	0.084	6.841	0.580	5.179	0.032	0.129
	Hájek	$-$ 0.078	2.125	$-$ 0.006	0.113	$-$ 0.611	2.281	$-$ 0.010	0.124	$-$ 0.001	0.107
	DR (outcome model misspecified)	0.029	1.650	0.109	1.308	$-$ 0.369	1.311	$-$ 0.019	0.920	0.103	1.310
	DR (outcome model correctly specified)	$-$ 0.001	0.109	$-$ 0.001	0.108	0.000	0.111	$-$ 0.001	0.109	$-$ 0.001	0.107

Note: The bias and RMSE are computed for Horvitz–Thompson (HT) estimator, the Hájek estimator, and the doubly robust (DR) estimator to compare the ATE estimation results. The true ATE is $10$ . The sample sizes are $n=200$ and $n=500$ . The number of simulations is $1000$ . GLM: generalized linear model; RMSE: root-mean-square error; ATE: average treatment effect; NNet: neural network; CBPS: covariate balancing propensity score.

Table 2.

Monte Carlo standard error (MCSE) for the performance measures of bias and RMSE. ⁵³

		GLM		GLM+IS		NNet		NNet+IS		CBPS
		Bias	RMSE	Bias	RMSE	Bias	RMSE	Bias	RMSE	Bias	RMSE
n = 200	PS model is misspecified
	HT	4.664	126.595	0.408	3.599	0.806	13.220	0.393	3.181	0.102	1.052
	Hájek	0.272	3.103	0.099	1.048	0.175	1.881	0.075	0.727	0.100	1.051
	DR (outcome model misspecified)	0.958	27.965	0.100	1.048	0.088	1.021	0.076	0.719	0.101	1.051
	DR (outcome model correctly specified)	0.012	0.239	0.006	0.038	0.007	0.073	0.006	0.039	0.006	0.038
	PS model is correctly specified
	HT	0.498	4.249	0.470	3.519	0.422	3.789	0.393	3.798	0.007	0.054
	Hájek	0.103	0.939	0.006	0.038	0.116	0.873	0.006	0.042	0.006	0.038
	DR (outcome model misspecified)	0.074	0.745	0.057	0.425	0.061	0.545	0.043	0.320	0.057	0.423
	DR (outcome model correctly specified)	0.006	0.039	0.006	0.038	0.006	0.043	0.006	0.040	0.006	0.038
n = 500	PS model is misspecified
	HT	11.025	329.136	0.177	1.356	1.960	59.302	0.163	1.339	0.070	0.908
	Hájek	0.334	4.550	0.069	0.904	0.169	3.121	0.046	0.471	0.069	0.907
	DR (outcome model misspecified)	2.918	87.252	0.070	0.910	0.111	2.492	0.046	0.467	0.070	0.906
	DR (outcome model correctly specified)	0.026	0.744	0.003	0.023	0.019	0.575	0.003	0.023	0.003	0.023
	PS model is correctly specified
	HT	0.293	2.254	0.259	1.764	0.216	2.225	0.163	1.306	0.004	0.029
	Hájek	0.067	0.608	0.004	0.024	0.070	0.513	0.004	0.027	0.003	0.023
	DR (outcome model misspecified)	0.052	0.552	0.041	0.322	0.040	0.393	0.029	0.219	0.041	0.317
	DR (outcome model correctly specified)	0.003	0.024	0.003	0.023	0.004	0.025	0.003	0.022	0.003	0.023

HT: Horvitz–Thompson; DR: doubly robust; PS: propensity score; GLM: generalized linear model; RMSE: root-mean-square error; ATE: average treatment effect; NNet: neural network; CBPS: covariate balancing propensity score.

Figure 3.

Boxplots of average treatment effect (ATE) estimates from $1000$ simulation runs for three ATE estimators based on different propensity score methods. Note: The data for the plot was truncated as the 0.1th and 99.9th percentiles of each estimator.

When the PS model was misspecified, the proposed calibration methods for PS estimation (NNet+IS and GLM+IS) generally led to a smaller bias and RMSE in all ATE estimators when compared to the corresponding standard estimation methods (NNet and GLM). When the PS model was correctly specified, the bias was small for all ATE estimators, but our proposed methods (NNet+IS and GLM+IS) had lower RMSE compared with the original (NNet and GLM) methods. Overall, the NNet+IS method performed the best, resulting in the low bias and lowest RMSE for the Hájek or the DR estimator. The impact of sample size on RMSE depended on the PS model. When the PS model was correctly specified, the RMSE decreased with the increasing sample size for all PS methods. However, when the PS model was misspecified, only NNet+IS consistently led to a smaller RMSE of ATE with increasing sample sizes.

Under the simulated scenario, HT estimators appeared to have much larger RMSE and bias than the other estimators. The improvement in the HT estimators was especially substantial when the robust PS estimation methods were used under the misspecified PS model. Compared to the CBPS method, our proposed methods, especially the NNet+IS, resulted in much smaller biases in the HT estimator although the RMSEs were higher. For the Hájek estimator, the NNet+IS method reduced both the bias and RMSE, and the GLM+IS method reduced the RMSE whether PS was correctly specified or not. NNet+IS had overall the best performance in terms of the lowest bias and RMSE compared to other PS estimation methods including the CBPS.

For the DR estimator, when neither the PS model nor the outcome model was correctly misspecified, our robust PS methods, GLM+IS and NNet+IS, led to a moderate to significant reduction in the RMSE when compared to the original GLM and NNet methods. In addition, the flexible neural network (NNet) model appeared to perform better than the simple logistic regression (GLM) since it was more robust to the model misspecification. When either the PS model or outcome model was misspecified, we still observed improvement in DR estimates using our robust PS estimation methods. The NNet+IS method outperformed all other methods including the CBPS method in terms of the bias and RSME of DR estimates. This confirmed that even with the doubly robust property, the DR estimator can still benefit from the robust PS estimation methods in finite samples. When both models were correctly specified, all the PS methods gave similar results as expected, and the DR estimator had a smaller bias and variance than HT and Hájek estimators that did not use outcome models. For the scenario when the PS model was correctly specified, our proposed PS methods (GLM+IS and NNet+IS) resulted in more improvement in the bias and RSME for the Hájek estimator, and the DR estimator under the misspecified outcome model, but less improvement for the HT estimator and DR estimator under the correctly specified outcome model, when compared to the corresponding original methods (GLM and NNet). On the other hand, the CBPS performed the best when used with the HT estimator. Interestingly, the Hájek estimator even outperformed the DR estimator with a misspecified outcome model when GLM+IS or NNet+IS method was used. It appeared that the misspecification of the outcome model may still negatively affect the DR estimator even though the PS model was correctly specified in this scenario.

Figure 4 shows the boxplots of standardized mean difference in each covariate (observed or latent) when the PS model was misspecified using $X^{*}$ . Each covariate was either unweighted or weighted by the inverse propensity scores estimated from different PS estimation methods. For the observed covariates $X_1^{*}$ to $X_4^{*}$ (Figure 4(a)), logistic regression (GLM) and neural networks (NNet) resulted in poor covariate balance while the proposed methods GLM+IS and NNet+IS showed significant improvements and CBPS had the best covariate balance. Note that the goal of the PS weighting method was to balance the latent covariates $X_1$ to $X_4$ , which were actually used to define the true PS model. When examining the standardized mean difference in these latent covariates (Figure 4(b)), we found that the NNet+IS achieved the best balance for all latent variables, followed by GLM+IS and CBPS.

Figure 4.

Propensity score (PS) weighted covariate balance under misspecified PS model. The sample size is $n=200$ : (a) covariate balance (observed covariates); and (b) covariate balance (latent covariates).

The mean absolute error (MAE), as a measurement of prediction performance of different PS estimation methods when the PS model is misspecified, was reported in Table 3. In this case, GLM and NNet have the smallest MAE of PS estimation when the sample size is 200 and 500, respectively. However, NNet+IS has the smallest bias of ATE estimation using the Hájek estimator regardless of the sample sizes as we shown in Table 2. These results were consistent with previous findings that the prediction error was not strongly related to the bias of ATE estimation.^2,10,16

Table 3.

MAE of estimated PS for different estimation methods when PS model is misspecified.

	GLM	GLM+IS	NNet	NNet+IS	CBPS
n = 200	0.081	0.088	0.088	0.095	0.089
n = 500	0.070	0.075	0.067	0.069	0.076

Note: Mean absolute error (MAE) is defined as MAE = $\frac{1}{n}\sum _i|{\hat{e}}_i-e_i|$ . PS: propensity score; GLM: generalized linear model; NNet: neural network; CBPS: covariate balance propensity score.

4. Data application: The right heart catheterization data

We demonstrated our proposed PS methods using the right heart catheterization (RHC) dataset from Connors et al. ⁵⁴ The objective of the analysis was to assess the effect of the utilization of RHC within the initial 24-hour period of care in the intensive care unit (ICU) on the hospital length of stay. The study cohort had 5734 patients in total, consisting of 2183 right heart catheterization users and 3551 non-users. The covariates identified by a panel of specialists in critical care included age, sex, race, years of education, income, medical insurance type, disease category, admission diagnosis, and other measurements.

In our analysis, we considered $65$ covariates for PS estimation. Two predictive base models, logistic regression and neural network, were used to estimate the PSs in the framework of the proposed calibrated loss function approach. We estimated the ATE of RHC with the HT, Hájek, and DR estimators and five PS estimation methods including GLM, GLM+IS, NNet, NNet+IS, and CBPS. We constructed the confidence intervals with $1000$ bootstrap sampling. For the DR method, we applied linear regression with the treatment assignment and all covariates in the outcome model. The comprehensive ATE estimation results were provided in Table 4.

Table 4.

Average effect of right heart catheterization on length of hospital stay (days).

	HT estimate of ATE	Hájek estimate of ATE	Doubly robust estimate of ATE
Methods for PS estimation	ATE (95% CI)	ATE (95% CI)	ATE (95% CI)
GLM	$2.49(0.56,4.41)$	$3.12(1.33,4.87)$	$2.86(1.04,4.54)$
GLM+IS	$2.40(0.05,4.08)$	$2.88(1.00,4.61)$	$2.86(0.98,4.60)$
NNet	$2.77(0.38,4.10)$	$3.66(1.69,5.43)$	$3.31(1.37,5.06)$
NNet+IS	$2.71(-0.24,3.92)$	$3.22(1.24,5.13)$	$3.21(1.22,5.05)$
CBPS	$2.92(1.02,4.66)$	$2.88(1.02,4.63)$	$2.88(1.01,4.62)$

HT: Horvitz-Thompson; ATE: average treatment effect; PS: propensity score; GLM: generalized linear model; NNet: neural network; CBPS: covariate balance propensity score.

Overall, the Hájek and DR estimates of ATE were nearly identical when using calibrated loss approaches (i.e. GLM+IS and NNet+IS) or CBPS for PS estimation. In contrast, the HT estimates were consistently lower than the Hájek and DR estimates, with the exception of when CBPS was used for the PS estimation. Furthermore, the DR methods, whether calibrated or not, produced similar estimates. The calibrated Hájek estimators yielded estimates that closely aligned with those of the DR estimates. Furthermore, the ATE estimates were larger when the neural network base model (NNet or NNet+IS) was used to estimate PSs.

As shown in our simulation results, when the outcome model was correctly specified, the DR estimator would be much less biased than the Hájek estimator. In practice, the PS model may not be perfectly specified. Given the minimal difference between these two estimates in our analysis, we also expected the outcome model to be less likely to be specified correctly. Therefore, we thought the results from the NNet+IS method were more reliable given this method demonstrated the best performance in our simulation study under misspecification of both the PS and outcome models.

We also examined the covariate balance achieved by four PS estimation methods and an unweighted one using the standardized mean difference defined by Austin et al.^13,55 (Figure 5). Most covariates can be better balanced via inverse PS weighting based on standard logistic regression and neural network methods and fall within the 0.1 threshold. ⁵⁵ However, the standardized difference in means for some covariates fell out of the $0.05$ threshold.^56,57 As we expected, the proposed calibration methods could further improve the covariate balance for these two prediction models and achieve almost exact balance.

Figure 5.

Standardized mean differences in inverse-propensity-score weighted covariate.

The dataset of Connors et al. ⁵⁴ has been analyzed with various approaches, while each employing a slightly different set of covariates. We summarized the reported effect sizes of RHC use and our estimates in Figure 6. Notice that some published results did not provide the confidence interval for the estimated effect size. Comparing our results with those reported previously, the overall conclusions for all methods were consistent. However, the proposed method using neural network base model resulted in a larger effect of RHC on the length of stays than most of these reported except the one presented by Li et al. ²²

Figure 6.

Forest plot of effect estimation given by multiple studies. Note: Keele and Small ⁵⁸ used targeted maximum likelihood (TMLE) with super learner and also generalized regression trees (GRF). ⁵⁸ Frank et al. ⁵⁹ applied TMLE. ⁵⁹ Li et al. ²² analyzed the right heart catheterization data with five methods including logistic regression, covariate balance propensity score (CBPS), covariate balancing scoring rules (CBSR), proposed propensity score with local balance 1 (PSLB 1), and proposed propensity score with local balance 2 (PSLB 2). ²²

5. Conclusions and discussions

The significance of the proposed methods lies in the ability to address critical challenges associated with PS-based causal effect estimation, particularly in the context of the PS model misspecifications. A correctly specified PS model theoretically ensures balanced covariate distributions between treatment groups.^2,3,60,61 Consequently, when the sample size is very large, the estimated PS could balance the covariates very well and the imbalance score penalty term is ignorable. When the sample size is limited, adding an imbalance penalty may correct the imbalance due to sampling variation so that increasing the precision. However, if the PS model is misspecified, the covariates may remain imbalanced even if the model accurately predicts treatment assignment. This misalignment can lead to biased estimates of the ATE. To address this issue, we developed robust PS estimation methods by calibrating the loss functions of predictive PS models. This adjustment considers the dual roles of PSs that the scores could predict the treatment and balance covariates. The proposed methods may improve upon prediction-focused methods when there is model misspecification due to complex, non-linear relationships among covariates. Our approach can mitigate bias and reduce RMSE in the estimation of the average treatment effect, particularly in scenarios involving model misspecifications.

As demonstrated in the simulation studies, when there was a model misspecification, prediction accuracy-based PS estimation methods such as logistic regression and neural network may result in large bias and variance in the ATE estimation. This problem was mainly driven by a small number of highly influential inverse probability weights that led to biased estimates and excessive variance. By adding a penalty term for covariates imbalance to the loss function, the proposed methods for PS estimation can achieve better covariate balance and eliminate the highly influential inverse probability weights, which led to a significant reduction in the bias and variance of ATE estimators. The proposed method also enables a better balance or even exact balance of observed covariates depending on the strength of the imbalance penalty. This type of balance shares similar spirits to the entropy balancing method ¹⁸ and the CBPS method, ²⁰ both of which achieve exact balance.

The DR methods have been developed to address potential misspecifications in either the PS or outcome models. These methods can provide consistent estimates even if one of the two models is misspecified, provided the other model is correctly specified. However, it was reported that the DR estimator may not always outperform a single outcome model when both the PS and outcome models are slightly misspecified. ¹¹ Through simulation studies (results not shown), we found that the DR estimator with NNet+IS as the PS model clearly outperformed the single misspecified outcome model, and the use of GLM+IS also showed improved performance. This highlights the ability of our proposed method in enhancing the DR estimator’s performance when both models are misspecified.

Our methods may help researchers avoid the ad hoc manipulation of estimated PSs or weights, such as trimming or truncating the weights. As shown in the simulation study, the PS estimation methods based on the calibrated loss function were more robust to model misspecification. In cases where either the PS model or the outcome model was correctly specified, the proposed calibrated version outperformed the prediction-based base model even for the DR estimator, especially in small sample sizes. When the PS model was correctly specified, the proposed methods could still reduce the variance in ATE estimation. In comparison to the logistic regression model, neural network model appeared to perform better under model misspecification. When the logistic regression was used as the prediction model for PS estimation, our proposed loss function calibration method and the CBPS method yielded very similar results for the Hájek and DR estimators. Both methods aim to optimize prediction and covariate balance at the same time. While the CBPS achieves covariate balancing within the framework of generalized method-of-moments, our approach is based on the calibrated loss function, which is more flexible to accommodate other covariate imbalance measures beyond marginal moments difference in each covariate. The penalty term to account for covariate imbalance can also vary across covariates depending on the strength of confounding effects. Lastly, our approach facilitates the use of other variable selection methods for predicting PS, such as least absolute shrinkage and selection operator (LASSO). ⁶² Specifically, when the dimension of covariates is high, an additional $l_1$ penalty can be incorporated to shrink the coefficients, thereby facilitating variable selection in conjunction with the proposed imbalance score penalty.

The proposed approach offers several advantages. First, it is robust to mild to moderate model misspecification. Even if the PS model is misspecified, the covariate balance is still controlled to a certain degree. Second, it is user-friendly, reducing the need for the back-and-forth covariate-balance check and model refit. Third, the framework is flexible in the sense that a variety of prediction models can be chosen, from the logistic regression model to machine learning methods such as neural networks. The only modification required is a simple calibration of the loss function. Lastly, the robust estimates of the PSs can be seamlessly integrated with a wide range of standard causal effect estimation methods, such as inverse probability weighting, matching, and stratification methods. Although the proposed method primarily focuses on enhancing PS estimation, integrating neural networks or other non-parametric approaches for outcome modeling may further enhance the performance of DR estimators. This integration allows for more flexible modeling, especially when the true outcome model formulation is unclear and likely to be misspecified with parametric regression models. ⁶³

In this article, we only consider a continuous outcome without censoring. However, extending our algorithms to other types of outcomes, such as binary and survival outcomes, is straightforward. Second, we investigated the most common scenario of a binary treatment in this work. Extensions to multi-level treatment groups warrant further study. While the proposed method is robust to misspecifications in the functional forms of the PS model, it still requires the absence of unmeasured confounding, as is the case with any observational study. Additionally, we apply an equal penalty to the imbalance across all covariates. In high-dimensional settings, confounder selection methods for causal inference are needed to help extend our proposed method to cases with high-dimensional covariates. ³² Moreover, achieving balance across all covariates may not always be feasible in high-dimensional settings. An outcome-assisted modeling approach motivated by Shortreed et al. ⁴³ may be considered to adjust the level of penalty for each covariate imbalance depending on the strength of the covariate-outcome association.

6. Implementation and software

The computational complexity of a neural network primarily depends on the architecture of the network. We run our simulation studies on Intel Xeon E5-2680 v3 with a 2.50 GHz CPU and 270 GB RAM. It roughly took <40 seconds under a sample size of 200, and <50 seconds under a sample size of 500 to implement NNet+IS method for a given tuning parameter with a single thread. The large tuning parameter selection range and CV will also increase the computational complexity. We used PyTorch’s autograd package that offers automatic differentiation for all tensor operations, enabling efficient gradient computation even with the imbalance penalty term in the loss function. Our code enables parallel computing to reduce the computational time, which is extremely helpful when running on high-performance computing machines. We implement the simulation and the real data applications in R 4.2.1 and Python 3.11.4. The source codes are available to the public via GitHub https://github.com/ys3298/RobustPSviaLossCalibration.