Propensity score synthetic augmentation matching using generative adversarial networks (PSSAM-GAN)

Abstract

Understanding causality is of crucial importance in biomedical sciences, where developing prediction models is insufficient because the models need to be actionable. However, data sources, such as electronic health records, are observational and often plagued with various types of biases, e.g. confounding. Although randomized controlled trials are the gold standard to estimate the causal effects of treatment interventions on health outcomes, they are not always possible. Propensity score matching (PSM) is a popular statistical technique for observational data that aims at balancing the characteristics of the population assigned either to a treatment or to a control group, making treatment assignment and outcome independent upon these characteristics. However, matching subjects can reduce the sample size. Inverse probability weighting (IPW) maintains the sample size, but extreme values can lead to instability. While PSM and IPW have been historically used in conjunction with linear regression, machine learning methods –including deep learning with propensity dropout– have been proposed to account for nonlinear treatment assignments. In this work, we propose a novel deep learning approach –the Propensity Score Synthetic Augmentation Matching using Generative Adversarial Networks (PSSAM-GAN)– that aims at keeping the sample size, without IPW, by generating synthetic matches. PSSAM-GAN can be used in conjunction with any other prediction method to estimate treatment effects. Experiments performed on both semi-synthetic (perinatal interventions) and real-world observational data (antibiotic treatments, and job interventions) show that the PSSAM-GAN approach effectively creates balanced datasets, relaxing the weighting/dropout needs for downstream methods, and providing competitive performance in effects estimation as compared to simple GAN and in conjunction with other deep counterfactual learning architectures, e.g. TARNet.

1. Introduction

In many research domains –especially in medicine and social science– it is challenging to distinguish conditional association from causation, yet such distinction is crucial for developing models that can be used to evaluate interventions, i.e. how actions can causally lead to desired effects. A typical causal inference scenario in medicine is to determine whether a treatment is effective with respect to an health outcome (e.g., lipid-lowering medications to reduce the risk of cardiovascular disease). The randomized controlled trial (RCT) randomly assigns treatments to individuals, making the treatment independent of individuals’ characteristics and guaranteeing that if an effect on an outcome is found, it can only be attributed to the treatment [23]. Nonetheless, performing RCT experiments is not always possible due to ethical concerns or operational grounds, e.g., it would not be ethical to assign and force individuals randomly to smoking for assessing if it caused lung cancer. Therefore, in many real-world situations, observational data are the only resource.

In observational studies, the knowledge on the data generation process or the space of variables involved in such process is limited. Hence, observational datasets can be contaminated with different biases, arising in various steps of the data generation process, and thus, impeding the way to infer causal claims [17]. Types of these biases include confounding (i.e., missing the true cause of an outcome but including spurious features correlated with the cause) and colliders (i.e., mistakenly including effects of an outcome as predictors), and other more complex configurations such as M-shaped or butterfly bias [3].

Therefore, inferring causal effects from observational data is problematic, and requires proper handling of possible biases on top of applying regression techniques. When analyzing the effect of a treatment intervention on a health outcome, propensity score matching (PSM) is often used to make the treatment independent from other pre-treatment covariates –in effect, balancing the existing data to resemble an RCT [18]. The propensity score π(x) represents the probability of receiving a treatment T = 1 (assuming that the alternative is no treatment T = 0) conditioned on the pre-treatment covariates X, denoted as

$$\pi (x)=P(T=1\mid \mathit{X}=\mathit{x})$$ (1)

The propensity score can be estimated using any regression method –often, regularized logistic regression is used [26]. Then, through PSM, the original dataset is re-balanced by matching pairs of treated/untreated individuals (i.e., cases vs. controls) by their propensity scores. Several algorithms can be used for the PSM process, such as k-nearest neighbor or caliper matching [2]. However, PSM often leads to exclusion of instances, and can evidently decrease the sample size or feature space. Thus, in real-world problems, (i) the number of controls often exceeds the number of cases, and (ii) the propensity score distribution of the case population can be seemingly shifted from the control population. Instead of PSM, other techniques such as inverse probability weighting (IPW) can be used, but they can also pose problems due to extreme values of the weights [6]. In general, PSM, IPW and related methods hinge upon the potential outcomes statistical framework [18, 19], which represents the observed (factual) and unobserved/potential (counterfactual) outcomes given a treatment, used to calculate causal treatment effects, both at the individual level and at the population (average) level.

Machine learning approaches have been used in the potential outcomes framework, e.g. Bayesian additive regression trees [7] and random forests [27]. The increasing availability of large amounts of electronic health record (EHR) also led to the flourishing of deep learning for causal inference [9], exploiting nonlinear approaches related to PSM and IPW. Notable examples include Treatment-Agnostic Representation Network (TARNet) [21], Dragonnet [22], and the Deep Counterfactual Network with Propensity-Dropout (DCN-PD) [1]. In TARNet, a two-heads multitask model was developed for estimating a binary treatment outcome, where each sample is assigned a specific weight indemnifying the imbalance between the treated and the control groups. During training, only parameters of one head of the network are updated, according to the treatment assignment indicator of the sample. Dragonnet was a modified TARNet which exploited targeted regularization using propensity scores. In DCN-PD, a doubly-robust model [4] was utilized, where a feed-forward network and a multitask network were developed to estimate the propensity score and the individual treatment effects, respectively. Then, a dropout probability of the network for each sample was calculated from its propensity score, and such dropout probability was used to regularize the network [25].

1.1. Contribution

In this work, we aim at addressing the problem of sample size reduction by PSM, without recurring to weighting schemes, to provide a general data re-balancing method applicable to any learning method downstream. We propose a novel deep learning approach –the Propensity Score Synthetic Augmentation Matching using Generative Adversarial Networks (PSSAM-GAN)– which leverages GAN [5] to generate data samples that can cover the treatment propensity space. In other words, PSSAM-GAN artificially creates treatment matches for unmatched controls (or vice versa). Deep generative models have been increasingly employed for causal inference problems, e.g. the Generative Adversarial Nets for Inference of Individualized Treatment Effects (GANITE) [29], Causal Effect Variational Autoencoder (CEVAE) [13], and Causal-GAN [11]. To the best of our knowledge, this is the first time GAN are not used directly for counterfactual prediction, but for separating the prediction task from the bias-handling. We here provide a generative model for data re-balancing, apt to aid not only with outcome prediction, but also with parameter inference.

After generating the unmatched samples with PSSAM-GAN, a semi-supervised method can be used to train a given multitask regression model, used to estimate average or individual treatment effects (through counterfactual prediction). In the experiments conducted here –using a semi-synthetic dataset from an RCT on perinatal intervention, and two real-world datasets on work interventions and antibiotic treatments– we show that PSSAM-GAN effectively creates balanced data, and favorably compares to the regularization approaches of DCN-PD and TARNet. In fact, PSSAM-GAN enables the application of a broader set of methods, not necessarily doubly-robust, guaranteeing the upstream handling of group assignment bias.

2. Methods

2.1. Problem formulation

We use the potential outcomes framework [18,19]. Thus, we consider a population sample of N individuals who can be prescribed a treatment T (binary, for simplicity), have a set of pre-treatment background covariates X, and have a measured health outcome Y. We denote each subject i as the tuple ${\{\mathit{X},T,Y\}}_{i=1}^N$. For each individual i the potential outcomes are represented as $Y_i^0$ and $Y_i^1$ when applying treatments T_i = 0 and T_i = 1, respectively. The individualized treatment effect (ITE) τ(x) for an individual i with feature vector X_i = x is defined as the difference in the mean potential outcomes under both treatment interventions (i.e., treated vs. not treated), conditional on the observed covariate vector x

$$\tau (\mathit{x})=\mathbb{E}\left[Y_i^1-Y_i^0\mid {\mathit{X}}_i=\mathit{x}\right]$$ (2)

Since an individual cannot be on and off treatment at the same time, i.e. one cannot obtain both potential outcomes from a person, we cannot calculate τ(x) directly. Only one outcome (factual) can be observed, while the other (counterfactual) is missing. If the potential outcomes are independent of the treatment assignment, conditionally on the background variables, i.e. {Y¹_, Y⁰} ⊥T|X, the assumption of strongly ignorable treatment assignment (SITA) is met [8,15]. Under the assumption of SITA, we can calculate ITE as follows: $\tau (\mathit{x})=\mathbb{E}\left[Y^1\mid T=1,\mathit{X}=\mathit{x}\right]-\mathbb{E}\left[Y^0\mid T=0,\mathit{X}=\mathit{x}\right]=\mathbb{E}[Y\mid T=1,\mathit{X}=\mathit{x}]-\mathbb{E}[Y\mid T=0,\mathit{X}=\mathit{x}]$.¹ Further, under SITA and by averaging over the distribution of X, we can calculate the average treatment effect (ATE) τ₀₁ as follows:

$${\tau }_{01}=\mathbb{E}[\tau (\mathit{X})]=\mathbb{E}[Y\mid T=1]-\mathbb{E}[Y\mid T=0]$$ (3)

ITE and ATE can be calculated with x being equally matched in treatment/control groups but stratification becomes infeasible as the dimension of x increases. With the help of PSM/IPW, through the conditional probability π(x) one can balance the probability of receiving T given X = x across the two comparison groups.

2.2. Proposed framework

The objective of PSSAM-GAN is to improve PSM by generating synthetic samples that can be paired to unmatched instances, covering the covariate space that would be otherwise lost, and re-balancing the dataset.

We outline the PSSAM-GAN procedure in Fig. 1. First, a regular PSM is performed on the original dataset. Second, a GAN is used to generate synthetic data, whose propensity scores would pair with the unmatched instances, leftover by PSM. Third, a semi-supervised learning approach assigns labels (i.e., outcomes) to the GAN-generated data. Finally, the synthetic data are collated to the original dataset and this balanced dataset can be used with regular learning algorithms.

Fig. 1.

Schematic of the PSSAM-GAN framework.

2.2.1. Propensity score matching

We calculate the propensity score vector π(X) for the dataset using a given regression method. Among the possible off-the-shelf choices, e.g. Bayesian additive regression trees, neural networks, logistic regression or LASSO, we used the neural network configuration used in the propensity network part of DCN-PD [1]. This choice gives more flexibility than logistic regression or LASSO in calculating propensity scores because it is able to address nonlinear or nonparallel treatment assignments. We then pair the treated and control samples according to their propensity scores. Here, we use a nearest-neighbor matching with replacement, which has been shown to provide a good bias-variance trade-off [2]. After matching, we pass the unpaired instances on to the GAN. For simplicity, we assume from now on that the unmatched instances are all controls (since with EHR usually the selection is in treatment assignment). Nonetheless, the GAN can generate samples to match also the leftover treated.

2.2.2. Generation of synthetic samples

Next, we use a GAN to generate a number of treated samples using the propensity score distribution of the unmatched controls as a guide. The GAN consists of two feed-forward networks, a generator and a discriminator, trained in the classical adversarial manner, such that the generator learns the true data distribution p(x) and fools the discriminator by producing fake samples. The standard optimization function to train a GAN with data distribution as p(x) and random noise as p(z) follows [5]:

$$\begin{gathered} \underset{G}{\min }\underset{D}{\max }V(D,G)={\mathbb{E}}_{\mathit{x}~p_{\text{data}}(\mathit{x})}[\log D(\mathit{x})] \\ +{\mathbb{E}}_{z~\mathcal{N}(\mathit{0},\mathit{1})}[\log (1-D(G(z))] \end{gathered}$$ (4)

The primary objective of the generator G in PSSAM-GAN is to learn the distribution of the unmatched control p_data(x_uc) such that it can fool the discriminator. To accomplish this, the generator is given a random noise $z~\mathcal{N}(\mathbf{0},\mathbf{1})$ as input. Along with this, an additional constraint to the generator is that the generated samples must have a similar propensity scores as the unmatched control ones. To achieve this, we introduce the following propensity loss that implements the propensity score constraint:

$${\mathcal{L}}_P^G\left(\pi \left({\mathit{x}}_{uc}^{(i)}\right),\pi \left(G\left(z^{(i)}\right)\right)\right)=(\pi \left({\mathit{x}}_{\mathit{u}\mathit{c}}^{(i)}\right)-\pi {\left(G\left(z^{(i)}\right)\right)}^2$$ (5)

where $\pi \left({\mathit{x}}_{uc}^{(i)}\right)$ and π(G(z⁽ⁱ⁾)) are the propensity scores of the i^th unmatched control and generated treated sample. As a result, G will output a full feature vector of a sample whose propensity scores will be as close as the unmatched control ones; at this point we set T=1 to consider them as treated. We show the architecture of the GAN block of PSSAM-GAN in Fig. 2. With the two objective functions of above, the generator G is iteratively optimized with m mini-batches as follows:

$$\begin{gathered} \underset{G}{\min }{\varphi }_g(D,G),\text{with} \\ {\varphi }_g(D,G)=\frac{1}{m}\sum _{i=1}^m[\log \left(1-D\left(G\left(z^{(i)}\right)\right)\right) \\ +\beta {\mathcal{L}}_P^G\left(\pi \left({\mathit{x}}_{\mathit{u}\mathit{c}}^{(i)}\right),\pi \left(G\left(z^{(i)}\right)\right)\right)] \end{gathered}$$ (6)

where ${\mathit{x}}_{uc}^{(i)}$ is the i^th unmatched control sample, and β > 0 is a hyperparameter. Similarly to GANITE [29], β has to be tuned during the experiments, to manage the relative weighting between the generator loss and the propensity loss.

Fig. 2.

Architecture of the generative adversarial network used in PSSAM-GAN to generate synthetic (treated) samples to be paired with the unmatched (control) samples by the propensity score matching.

The discriminator D in PSSAM-GAN is a standard discriminator whose primary objective is to differentiate the real unmatched control and fake generated samples. The discriminator D is optimized with m mini-batches as follows:

$$\begin{gathered} \underset{D_G}{\max }{\varphi }_d(D,G),\text{with} \\ {\varphi }_d(D,G)=\frac{1}{m}\sum _{i=1}^m[\log D\left({\mathit{x}}_{\mathit{u}\mathit{c}}^{(i)}\right) \\ +\log \left(1-D\left(G\left(z^{(i)}\right)\right)\right] \end{gathered}$$ (7)

2.2.3. Synthetic labeling

The GAN generates unlabeled instances. In principle, outcome assignment based on nearest-neighbor match could be used, but rather than the sole comparison of π(X), one should match on X, thus needing a specific distance function. Another, more general-purpose option is semi-supervised learning [30]. Here, we employed multitask regression by TARNet and DCN-PD, using the regularizer, to infer the outcome labels Y for the generated samples. In this way, we were able to obtain both a factual and counterfactual outcomes, and then use either one based on treatment assignment. Of note, we tested the label assignment using the original dataset and also only the PSM-induced subset.

2.2.4. Downstream model learning

The original dataset is merged with the synthetic one into a balanced final dataset, which is used to train a multitask regression model such as TARNet or DCN-PD. Note that when creating the augmented dataset, some duplicates from the initial PSM stage can be included. Models are then evaluated in terms of ITE/ATE using a holdout test set which was never used in any previous step.

Algorithm 1 describes the different phases of PSSAM-GAN algorithm to obtain the final model. Note that by removing the propensity loss in Eq. 5, it is possible to obtain a simple GAN, consistent with the one proposed in [5]. The simple GAN will be compared with the proposed PSSAM-GAN in the experimental settings, showing the performance gain provided by incorporating the propensity loss equation.

2.3. Experimental setup

2.3.1. Datasets

Evaluating causal inference algorithm on real-world data is always difficult because of the problem of underlying unmeasured bias and missing counterfactual outcomes. Nonetheless, these data are crucial for evidence-based validation. In this work, we used both semi-synthetic (i. e., real factual outcomes and simulated counterfactuals) and real data.

The semi-synthetic dataset is derived from the Infant Health and Development Program (IHDP), a real RCT for improving health outcomes of premature infants (https://www.icpsr.umich.edu/web/HMCA/studies/9795). The original data was artificially modified to induce bias, by selecting only one ethnic group (non-white) in one intervention arm, while the other included all. Then, a nonlinear surface outcome function was designed to fit a given treatment effect, providing counterfactual outcomes. The final dataset consists of 747 subjects –139 treated and 608 controls– with 25 covariates associated with each subject [7].

Next, we use the Jobs dataset [12,24], which includes two population samples subject to work interventions and income survey: the LaLonde’s National Supported Work Demonstration (NSWD) experimental sample (297 treated, 425 control) and the Population Survey of Income Dynamics (PSID) comparison group (2,490 control). This is real data that does not involve any alteration or simulation. There were 482 (15%) subjects unemployed by the end of the study. The treatment variable was job training and the outcomes were post-training income and employment status. A binary classification task was set up as described in [21]. Both IHDP and Jobs have been used widely as benchmarks in causal inference studies (https://www.fredjo.com/).

Algorithm 1.

PSSAM-GAN: Propensity Score Synthetic Augmentation Matching using Generative Adversarial Networks.

Input: Training set X = {(x(1), t(1), y(1)),…, (x(n), t(n), y(n)) }, loss function L(., .), hyperparameter β > 0, generator network G with initial parameters θg, discriminator network D with initial parameters θd, epochs K

1. Perform neural network-based propensity score π(X) calculation and matching between the treated samples (t = 1) and control samples (t = 0) using the nearest neighbor algorithm with replacement

2. Separate out the unmatched control samples from the matched ones

3. for epochs = 1, 2, … ,Ks do

4.   Sample minibatch of m noise samples ${\left\{z^{(i)}\right\}}_{i=1}^m$ where $z~\mathcal{N}(0,1)$

5.   Sample minibatch of m unmatched control samples from the distribution pdata(xuc)

6.   Update the discriminator D by ascending its stochastic gradient: ${\nabla }_{{\theta }_d}\left({\varphi }_d(D,G)\right)$

7.   Sample minibatch of m noise samples ${\left\{z^{(i)}\right\}}_{i=1}^m$ where $z~\mathcal{N}(0,1)$

8.   Update the generator G by descending its stochastic gradient: ${\nabla }_{{\theta }_d}\left({\varphi }_d(D,G)\right)$

9. end for

10. Remove the discriminator D and get the generated samples from the generator G and assign them as treated

11. Train a multitask regression model with the original training set X to obtain outcome labels for the generated treated samples

12. Merge the original dataset with the synthetic data into an augmented balanced dataset

13. Train a multitask regression model with the augmented dataset

Next, we use a collection of bacterial samples (Staphylococcus aureus) sequenced by Multilocus sequence typing (MLST) with associated information on antibiotic resistance (https://pubmlst.org/), which was previously used to identify genetic resistance signatures [16]. This is a novel, real dataset. Bacteria acquire resistance to antibiotics through various genetic mechanisms, e.g. mutations or horizontal gene transfer, but configurations for housekeeping genes are usually limited to few, specific point mutations. Therefore, the propensities of treatment resistance or susceptibility for most mutations in housekeeping genes should be similar. However, much larger discriminative value of housekeeping mutations was found, which could be explained by selection bias [16]. The S. aureus dataset included 2006 samples, over 3000 housekeeping mutation variables, and the methicillin antibiotic treatment.

2.3.2. Deep learning modules’ settings

For the GAN module, we use a 3-layer feed-forward neural network for both the generator and discriminator with LeakyRelU [28] activations in the hidden layers. The generator includes between 17 and 25 hidden units in each layer (tuned for each dataset), whereas the discriminator included 25 hidden unites. The generator/discriminator learning rates and the hyperparameter β were set to 0.0002 and 1.0, respectively. The GAN was trained for 10,000 epochs (K_s) with a mini-batch size (m) of 64. The parameters of the generator (θ_g) and discriminator (θ_d) were updated by the Adam optimizer [10].

We evaluated PSSAM-GAN under different multitask regression configurations, varying weighting and dropout schemes: (i) DCN-PD; (ii) DCN without PD; (iii) DCN with constant dropout probability of 0.5 for all layers and all training examples; (iv) TARNet; and (v) the GAN without the propensity loss (defined by Eq. 5). The methods were trained both with the original dataset and with the augmented balanced dataset which included the synthetic instances. DCN-PD employs an alternate training strategy where either treated or control samples are used when the epoch number is even or odd, respectively. In TARNet, each sample is weighted, and the mini-batches are sampled randomly [21].

We ran the DCN and PD network for 400 and 50 epochs, respectively, in both the semi-supervised step and the final regression module. The learning rates of the DCN and PD were 0.0001 and 0.001 respectively. For TARNet, the model was trained for 3000 epochs in both the semi-supervised learning and final regression. We set the learning rate, the λ L2 regularizer (for the parameters of hypotheses layers), and the batch size to 0.0001, 0.00001, and 100, respectively. We used the default parameters and training settings for DCN-PD [1] and TARNet [21]. PSSAM-GAN was implemented using Pytorch (https://pytorch.org/). All the experiments were performed on an i7 Macbook Pro, with 16GB RAM, mounting Catalina OS. The code is available under the MIT license on GitHub at: https://github.com/Shantanu48114860/PSSAM-GAN.

2.3.3. Validation and performance metrics

When comparing PSSAM-GAN with DCN-PD and TARNet for the Jobs dataset, we followed the original papers’ train/test split setup for continuity in interpreting the results, noting that some of the original procedures have been re-coded here in Pytorch. For the IHDP dataset, we trained the models on 80% of the data and tested on 20%, as described in [1], and repeated this for 1000 iterations. For the Jobs dataset, we selected 10% of the data as the test set, and repeated the procedure for 10 iterations [21].

In terms of performance metrics, for IHDP – according to [7,21,29] – using the expected Precision in Estimation of Heterogeneous Effect (ϵ_PEHE) and the ATE, we evaluated $\sqrt{{ϵ}_{PEHE}}$ and the error on the true ATE, i.e. ϵ_ATE. Unlike IHDP, there are no counterfactual outcomes in the Jobs dataset, so the performance metrics used were Policy Risk (R_pol(π)) and ϵ_ATT, which is the error on the ATE in the treated group. The detailed formulae are given as follows:

$$\begin{gathered} {ϵ}_{PEHE}=\frac{1}{N}\sum _{n=0}^N({\mathbb{E}}_{y_1(n),y_0(n)~{\mu }_Y(\mathit{x}(n))}[y_1(n)-y_0(n)] \\ {-\left[\widehat{y_1}(n)-\widehat{y_0}(n)\right])}^2 \end{gathered}$$ (8)

$${ϵ}_{ATE}={\Vert \frac{1}{N}\sum _{n=0}^N{\mathbb{E}}_{y(n)~\mu (n)}[y(n)]-\frac{1}{N}\sum _{n=0}^N\widehat{y}(n)\Vert }_2^2$$ (9)

$$\begin{gathered} R_{pol}(\pi )=\frac{1}{N}\sum _{n=0}^N[1-(\sum _{i=1}^k[\frac{1}{\left|{\mathrm{\Pi }}_i\cap T_i\cap E\right|} \\ \sum _{\mathit{x}(n)\in {\mathrm{\Pi }}_i\cap T_i\cap E}{y}_i(n)\times \frac{\left|{\mathrm{\Pi }}_n\cap E\right|}{\left|E\right|}])] \end{gathered}$$ (10)

where ${\mathrm{\Pi }}_i=\{\mathit{x}(\mathit{n}):i=\arg \max \widehat{y}\}$,

T_i = {x(n) : t_i(n) = 1}, and E is the subset of RCT.

3. Results

The GAN module produced synthetic instances for the unmatched controls covering the whole propensity score space in both datasets. Fig. 3 shows the distributions of the propensity scores, stratified by treatment group, for the original dataset, the PSM-filtered data (only treated samples matched to their controls using nearest neighbor), and then the final dataset that merged the GAN’s synthetic samples with the original data. The original score distributions are highly imbalanced and show substantial distribution shift between treated and control samples; the PSM-filtered distributions are balanced, but included a smaller sample size over a narrower score interval. Instead, the GAN-augmented dataset not only covers a larger propensity score interval, but also show that the stratified treated/controls distributions are very similar.

Fig. 3.

Histograms of the propensity score distributions for IHDP (top) and Jobs (bottom) datasets, stratified by treatment/control group. From left to right, the panels show the original dataset before PSM, the matched dataset after PSM and the dataset augmented with PSSAM-GAN framework.

Table 1 shows the mean ± st.dev performance metrics on both IHDP and Jobs dataset when applying the PSSAM-GAN framework and comparing it with the regular weighting/dropout deep learning models. Note that the point estimates for DCN(-PD) and TARNet may be slightly different from the referenced papers because we re-implemented some of the procedural steps using Pytorch. Nonetheless, the confidence intervals are consistent with prior literature. PSSAM-GAN improves performance over the other approaches, including the simple GAN, and results show that best and near-best results can be achieved without weighting or dropout regularization schemes. The advantage of PSSAM-GAN over the simple GAN is evident across both the IHDP and Jobs datasets, and consistent with any of the downstream methods used. This is due to the incorporation of the propensity loss in Eq. 5, which ensures that the synthetically generated treated samples have propensity scores closer to the control samples, including those who were unmatched. In addition, PSSAM-GAN had superior performance over DCN-PD and TARNet by 33% (without applying PD regularization) and 5%, respectively, in the IHDP dataset, whilst showed a more moderate improvement in the Jobs (1–2% for both performance indices). Of note, the moderate improvement in the Jobs dataset is not unexpected, as it was previously shown that the average gain of TARNet over ordinary least square was minimal (not significant) for the policy risk (R_pol(π)) and there was no gain for ϵ_ATT [21]. In our experiments, we found that the increment in performance obtained with PSSAM-GAN was highly unlikely due to chance, as in all but one tests the null hypothesis of no difference was not supported [14]. The p-values reported in Table 1 refer to ϵ_ATE and ϵ_ATT, while for the other two measures they were all below 0.0001.

Table 1.

Performance on the out-of-sample test sets (mean ± st.dev) of the PSSAM-GAN framework vs. a simple GAN setup, and the regular weighting/dropout deep learning models, on the IHDP and Jobs datasets.

	IHDP			Jobs
	$\sqrt{{ϵ}_{PEHE}}$	ϵATE	p-value	ϵATT	p-value	Rpol(π)
DCN-PD	1.77 ± 0.09	0.18 ± 0.18	0.3294	0.08 ± 0.07	0.0192	0.29 ± 0.05
DCN (Dropout probability: 0.5)	2.08 ± 0.12	0.43 ± 0.17	<0.0001	0.11 ± 0.05	<0.0001	0.29 ± 0.05
TARNet	1.13 ± 0.03	0.21 ± 0.13	<0.0001	0.13 ± 0.07	<0.0001	0.29 ± 0.06
PSSAM-GAN
+ DCN	1.18 ± 0.03	0.19 ± 0.09	0.0041	0.07 ±0.04	ref.	0.29 ± 0.05
+ DCN-PD	1.82 ± 0.14	0.52 ± 0.38	<0.0001	0.07 ± 0.07	1.0000	0.29 ± 0.05
+ DCN (Dropout probability: 0.5)	2.07 ± 0.12	0.45 ± 0.21	<0.0001	0.09 ± 0.08	<0.0001	0.29 ± 0.06
+ TARNet	1.07 ±0.02	0.17 ±0.10	ref.	0.13 ± 0.10	<0.0001	0.27 ±0.03
GAN
+ DCN	1.43 ± 0.04	0.19 ± 0.1	0.0064	0.10 ± 0.07	<0.0001	0.29 ± 0.05
+ DCN-PD	1.98 ± 0.16	0.82 ± 0.44	<0.0001	0.08 ± 0.07	0.0192	0.30 ± 0.05
+ DCN (Dropout probability: 0.5)	2.21 ± 0.13	0.44 ± 0.16	<0.0192	0.12 ± 0.07	<0.0001	0.32 ± 0.05
+ TARNet	1.37 ± 0.02	0.23 ± 0.09	<0.0001	0.13 ± 0.17	<0.0001	0.30 ± 0.05

Next, we analyzed the larger bacterial dataset of MLST with respect to antibiotic treatment resistance. Fig. 4 shows the propensity score distributions for methicillin in the original data, in the matched subsets by PSM, and in the augmented space by PSSAM-GAN. As expected, the PSM drops a consistent part of the population, methicillin-susceptible in this case. On the contrary, PSSAM-GAN generates examples that cover the whole spectrum of treatment propensities, yet some of the high scores for methicillin resistance are left unmatched. In fact, we cannot rule out that some of the mutations in the genes sequenced by MLST might be actually causative of methicillin resistance. However, for most of the housekeeping mutations (and here we are using over 3000 sites) this would not be the case, and rather the association with treatment resistance would be due to founder events and hitchhiking mutations, as suggested in a previous study on the same data looking at phylogenetic structure and sampling bias [16]. In fact, the MEGARes database (https://megares.meglab.org/), one of the most comprehensive collections of bacterial genes with known association to antibiotic resistance, lists only 490 single-point mutations of housekeeping genes with confirmed resistance compared to over 7800 resistance entries of other genes across all antibiotics (i.e. 6.2% of known resistance is found in housekeeping genes). On average, housekeeping genes in MEGARes have a median (interquartile range) nucleotide length of 1544 (802–2, 619), thus the expected contribution of single-point mutations to any-antibiotic resistance is only 0.06%.

Fig. 4.

Histograms of the propensity score distributions for the bacterial dataset. From left to right, the panels show the original dataset before PSM, the matched dataset after PSM and the dataset augmented with PSSAM-GAN framework.

4. Discussion

We show that PSSAM-GAN is able to generate synthetic instances that cover the space of propensity scores for unmatched samples, increasing the sample size and richness of the covariate space. In addition, we demonstrate that it compares and combines favorably with the deep learning weighting/dropout schemes of DCN-PD and TARNet (besides being better than a baseline GAN architecture). We note that the provision of balanced datasets offers an operational advantage since other multitask regression learning can be applied without the need to tailor them to handle group assignment imbalance. For instance, PSSAM-GAN could be further integrated with Dragonnet [22].

The utility of PSSAM-GAN in practice can be foreseen in various scenarios, one such being its employment in emulating RCT on observational data, i.e., the target trial framework. The inclusion criteria apt to fulfill the SITA in target trial designs often lead to a case funnel, with substantial drop in sample size from the initial population (which can be large) and possibly imbalanced intervention groups.

Limitations of this contribution include the possible performance fluctuations due to the initial choice of the PSM method and propensity score function. Although propensity-based methods are less problematic than IPW because they do not have to invert weights, skewness of their distribution can still affect model learning and generation of samples. Similarly, performance can vary on the basis of the dropout threshold, as our results showed –error increased using a dropout probability of 0.5. Possible improvement could be achieved by using stable weights [31]. Another possible issue is that the semi-supervised step could be biased by data imbalance, and that the labelling might be inaccurate. Yet, the label assignment step is a pure prediction task and we are not interested in calculating any effect at that point of the procedure; in fact, training with only matched samples (to further reduce bias) did not affect the overall performance.

Here, we explored PSSAM-GAN in the special case of a binary treatment/intervention. One extension to this work is developing and applying it to multiple treatment groups, e.g. using a softmax. Also, in the downstream learning module, when using TARNet, one could generate balanced mini-batches using the perfect-match algorithm, leveraging the augmented-balanced dataset, instead of the weighting scheme [20]. Finally, for high-dimensional datasets, a regularization scheme or dimension reduction (e.g. through a latent space from a sparse autoencoder) could be used to reduce noise.

4.1. Conclusion

In conclusion, PSSAM-GAN is a promising approach to overcome the limitations of PSM- and IPW-based methods, especially addressing the sample size reduction and the weight instability problems. Differently from other approaches, PSSAM-GAN not only provides treatment effect estimation, but also sample generation. PSSAM-GAN is also well suited for high-dimensional studies, e.g. using EHR or genomics. Further, it provides a general-purpose tool apt to be used in conjunction with any multitask regression scheme.