Predictive State Propensity Subclassiﬁcation (PSPS): A causal inference algorithm for data-driven propensity score stratiﬁcation
Abstract
We introduce Predictive State Propensity Subclassiﬁcation (PSPS), a novel learning algorithm for causal inference from observational data.
PSPS combines propensity and outcome models into one encompassing probabilistic framework, which can be jointly estimated using maximum likelihood or Bayesian inference.
The methodology applies to both discrete and continuous treatments and can estimate unit-level and population-level average treatment effects.
We describe the neural network architecture and its TensorFlow implementation for likelihood optimization.
Finally we demonstrate via large-scale simulations that PSPS outperforms state-of-the-art algorithms – both on bias for average treatment effects (ATEs) and RMSE for unit-level treatment effects (UTEs).
1. Introduction
Measuring the impact of an intervention on an outcome of interest is ubiquitous across science, business, and engineering.
For example, researchers want to ﬁnd out whether a new drug has the desired effect on blood pressure reduction; advertisers would like to know whether an ad worked to increase sales; or a software company wants to evaluate the impact of a new feature launch on improved user experience.
In an ideal world a randomized experiment would be run to estimate the impact of each of these interventions.
Randomization ensures balance in expectation on important confounders – variables that affect both who got treated and the outcome – across the different treatment arms allowing for unbiased estimates of the causal effect of the intervention.
An experimenter may furthermore block or stratify on certain features in order to ensure balance across treatment assignments for a particular randomization (Bernstein, 1927).
Randomized experiments, however, are not always feasible due to time, cost, legal, or ethical reasons – to name a few – in which case the decision maker has to rely on observational data.
Their lack of randomization leads to feature imbalances between the treatment arms and thus make identiﬁcation and estimation of causal effects difﬁcult.
Causal inference algorithms for observational data cannot rely on randomization, but have to use other means to mitigate the imbalances.
As such they are often viewed as tools to help ﬁx broken randomized experiments (Rubin, 2008), by trying to recreate balance between different treatment levels that did not occur due to the lack of randomization.
It is important to emphasize that in order to make good data-driven decisions on how to intervene in real world settings, algorithms need to provide unbiased and low-variance causal effect estimates.
Learning algorithms that solely optimize for predictive accuracy are likely diametrical to accurate causal effect estimates as they fall victim to selection bias & confounding (Bareinboim et al., 2014).
In this work we introduce Predictive State Propensity Subclassiﬁcation (PSPS), a novel causal learning algorithm that leverages multi-task learning (Ruder, 2017).
PSPS follows the Rubin Causal
Model (RCM) framework (Imbens and Rubin, 2010) and is in the same vein as matching, IPW and traditional subclassiﬁcation that aim to ﬁx the broken design of the experiment.
Our key contributions are three-fold: a) we incorporate predictive state smoothing (PRESS) (Goerg, 2017, 2018) for the propensity score, which yields a principled data-driven way to obtain the strata or blocks used in traditional matching approaches; b) the entire causal inference algorithm can be represented as an end-to-end multi-task neural network, which allows for – statistically and computationally – efﬁcient training and inference through a single computational graph; c) we derive a causal loss function that balances predictive performance with properly accounting for selection bias.
1.1. Literature Review
Existing methods that aim to ﬁx non-randomization in observational data are often based on matching (Stuart, 2010) or inverse propensity weighting (IPW) coupled with outcome modeling; e.g., entropy balancing (Hainmueller, 2012) and double robust estimation (Robins and Rotnitzky, 1995;
Hirano et al., 2003; Lunceford and Davidian, 2004).
Subclassiﬁcation addresses feature imbalance by creating strata of observations such that within each stratum there is approximate feature balance between treated and control units (Cochran, 1968).
We do not use do -calculus (Pearl, 1995) in this work, but since PSPS is a fully probabilistic framework it is straightforward to use plug-in causal predictions and conditional and joint distribution estimates to a do -calculus based causal analysis.
Recently causal inference has also received attention in the machine learning community: for example via counterfactual representation learning (Johansson et al., 2016) or Dragonnet (Shi et al., 2019).
Section 4 compares PSPS to various methods on causal benchmark datasets.
2. Methodology, Neural Network Architecture & Learning Objectives
Let T be the treatment variable, Y the observed outcome, and X is a collection of pre-treatment features.
For example, the canonical causal estimation scenario is the case of binary treatment,
T ∈ { 0 , 1 } , with continuous outcome Y ∈ R , and k -dimensional feature space X ∈ R k .
We will revisit an example of continuous treatment in Sections 3.
Unless speciﬁed otherwise though, we intentionally leave the dimensionality and variable types of ( T , Y , X ) unspeciﬁed as the core concepts and derivations also hold for multivariate outcomes and multi-level or continuous treatments.
2.1. Factorization of the Joint Distribution of Treatment & Outcome
Statistically, treatment assignment and application are confounded when learning from non-randomized data.
In essence many causal methods aim to break this confounding between assignment and the application effect of treatment on the outcome.
In PSPS this dual role is made explicit by
[Equation]
[Figure]
where θ Y and θ T parameterize the conditional distributions of Y and T , respectively.
The factorization in (1) makes it explicit that T plays two roles: i) once as the output of a propensity model (treatment assignment, T assign ), and ii) as an input to an outcome model (treatment application, T apply ).
Since assignment and application are confounded we modulate the distribution on the observed outcome ( Y | T apply , X ) with the fact that there was self-selection to treatment ( T assign | X ).
This distinction also illustrates that practicioners are truly interested in the effect of applying treatment t on unit i with observed features x i , i.e. p θ Y
( Y i | T apply =
t , X = x i ) .
While treatment assignment and application are the same random variable and data, respectively, we keep the distinction for clarity on the computational graph.
Figure 1 shows that T assign and T apply play two different roles: T assign is an output while T apply is an input.
It is important to have separate nodes – rather than a single “ T ” node – as this would result in an incorrect computational representation, and hence implementation.
2.2. Propensity Score Subclassiﬁcation
Under a randomized-block experiment the propensity score P ( T | X ) is known by construction and the properties of randomization meet the ignorable treatment assignment (“ignorability”) assumption (Rosenbaum and Rubin, 1983) allowing for the identiﬁcation and estimation of causal effects.
For observational data studies treatment assignment is typically unknown and as such the ignorability assumption is usually made on a leap of faith after adjusting for any observed imbalances.
To deal with potential self-selection bias and imbalances between treated and control units,
Rosenbaum and Rubin (1985) proposed subclassiﬁcation , which stratiﬁes units by a coarsened value of predicted propensity scores ˆ P ( T | X ) , – usually achieved by binning; it is also often accompanied by a model-based adjustment to deal with any remaining imbalances.
The idea being that observations within each bin have approximately the same propensity score and thus are more robust to any misspeciﬁcation in the outcome models of each bin.
The (predicted) outcome for treated and controls can then be compared within each stratum to estimate the effect of treatment.
One common problem with propensity score subclassiﬁcation is that there are little to no guidelines on how to actually create the subclasses.
Common heuristics are to a) bin on quantiles, b) have equally sized bins, or c) try to at least make the largest bin size moderately small.
Imbens and
Rubin (2015) create bins recursively by testing equality of propensity score means between treated and control via a t-test.
If the difference is signiﬁcant then the stratum is split at the median and the procedure repeats on sub-strata until the test statistic is not signiﬁcant or a pre-speciﬁed minimum number of treated and controls are met.
The end result, however, is inevitably an ad-hoc heuristic.
2.3. Predictive States As The Optimal Data-Driven Subclassiﬁcation
PSPS follows the same experimental design as traditional subclassiﬁcation except for how subclasses are created and how subsequent analyses take place.
The core causal assumptions such as “ignorability” are the same and we rely on these established results (Imbens and Rubin, 2015).
PSPS puts the goal of constructing subclasses front and center by learning the optimal stratiﬁcation directly from the observed data using Predictive State Smoothing (PRESS) (Goerg, 2017, 2018).
PRESS represents the conditional predictive distribution p ( T | X ) as a mixture over J latent predictive states, { s 1 ( X ) ,..., s J ( X ) }
(Figure 1).
Each state is an equivalence class over distribution space such that observations with the same conditional distribution are in the same predictive state.
That is X i
and X i 2 have the same state s j if and only if
[Equation]
The key insight here is that predictive states are exactly the strata { s 1 ,..., s J } that would have ideally been blocked on when assigning treatment: knowing the predictive state of observation i – which is unobserved in observational data –would be sufﬁcient to knowing the propensity score distribution of all observations in that state.
Predictive states are also minimal sufﬁcient for prediction, i.e., they make T conditionally independent of X given s j (Goerg, 2017).
Hence
[Equation]
where w j ( X ) :
= p ( S = s j | X ) denote the probability that X belongs to state s j .
We use the weight-vector notation w = ( w 1 ,..., w J ) as the representation of the state space mapping; by deﬁnition ∑ Jj = 1 w j = 1 and w j ≥ 0.
It is easy to see that two observations i
and i 2 – with potentially different features X i 1 and X i 2 – that are mapped to the same predictive state with same probabilities, i.e., w i 1 ≡ w i 2 , necessarily have the same treatment distribution since p ( T i |
X i ) =
∑ J j = 1 w i , j · p ( T i | s j ) .
Put in other words, if i 1 and i 2 have different observed features, x i 1 and x i 2 , but are in the same state, s j , then they have the same propensity score:
[Equation]
This is exactly the underlying goal of traditional subclassiﬁcation by Imbens and Rubin (2015).
Another way of thinking about it is that each state is a matched group where matching occurs based on an optimal distance between units i
and i 2 with features
X
i {+1 and X i 2 measured as the distance between propensity score distributions P ( T i 1 | X i 1 ) and P ( T i 
 | X i 2 ) .+}
2.4. Combining Outcome Models Across Predictive States
After constructing the subclasses, PSPS follows the traditional subclassifcation philosophy of training an outcome model in each subclass (Figure 1).
This has the effect of adjusting for features if the PRESS propensity model is misspeciﬁed.
Even when it is correctly speciﬁed, more accurate outcome models can further reduce the variability in causal effect estimates (Rosenbaum, 2002).
To replicate the behavior of estimating the causal effect within each subclass we use the propensity model predictive states to factorize the conditional outcome distribution in (1) as
[Equation]
(5) where the lift ratio ρ i , j ∈ [ 0 , ∞ ) quantiﬁes how much the state predictive distribution, p ( T i | s j ) , differs from the point prediction of the propensity score for feature x
i , p ( T i | X i = x i ) , and θ ( j ) Y parameterizes the outcome model in state j .
Since v i , j : =
ρ i , j · w i , j = P ( S = s j | X i , T i ) is a proper probability mass function for each i , the decomposition in (5) is a mixture over state-conditional outcome models, p θ ( j ) Y ( Y | T , X ) , where model j contributes with weight v i , j to the ﬁnal estimate for unit i .
That means observation i is only partially contributing to state j and also explains why, in general, outcome model parameter estimates will be different for every state.
Consider the special case where each unit i is mapped exactly to one state with probability(i.e.
, a deterministic mapping), then also ρ i , j = 1 for exactly one j and 0 otherwise.
In that case,
[Equation]
which is equivalent to binning propensity scores and training an outcome model in each bin.
Thus,
PSPS is a probabilistic generalization of traditional – and often heuristic – propensity score binning.
2.5. Model Training via Multi-Task Learning
Since PSPS provides a consistent causal probabilistic framework, we use maximum likelihood estimation (MLE) as the primary learning strategy; we leave Bayesian inference for future work.
Decomposing the joint distribution in (1) allows us to specify propensity model and outcome models separately, yet link them through their joint log-likelihood function as
[Equation]
where (cid:96) ( τ ; Z ) :
= log p τ ( Z ) is the log-likelihood of τ for data Z .
This decomposition, also used for Targeted MLE (TLME)
(Van der Laan and Rose, 2011), naturally leads to framing parameter estimation as a multi-task learning problem (Ruder, 2017) with a negative log-likelihood loss
L ( [Equation] joint loss
[Equation]
outcome model loss + propensity model loss [Equation] where ˆ T assign ( θ T ; X ) is the predicted propensity of assigning treatment and ˆ Y ( θ Y , θ T ; X , T apply ) is the predicted outcome when applying treatment.
The MLE can be obtained by solving
[Equation]
For randomized experiments the second term (8) is constant since T is independent X ; hence
PSPS reduces to a simple outcome model when faced with properly randomized treatment.
2.6. Separation of Design and Analysis
In observational studies, practitioners often face a choice between only training an outcome model ( P ( Y | T , X ) ) or a two-step procedure of training a propensity model
ﬁrst ( P ( T | X ) ) and an adjusted/weighted outcome model later.
PSPS includes both approaches as special cases, and allows practitioners to gradually move between them using a tunable hyperparameter.
The decomposition in (8) can be seen as trying to ﬁnd an optimal outcome model, but with a penalty for bad propensity models.
It is then natural to introduce a penalty parameter, λ ∈ [ 0 , ∞ ) ,
[Equation]
which reduces to (8) when λ = 1.
This can be interpreted as just training a single outcome model
Y | T , X , with a penalty L ( T , ˆ T ( θ T )) that adjusts for feature distribution imbalances between treated and controls (see also counterfactual representation learning in Johansson et al., 2016).
When λ =
the second term vanishes from (10) and the optimizer only focuses on outcome models, weighted by a random mixture of predictive state weights from each state, which – in expectation – is just a simple overall outcome model.
For λ → ∞ the optimizer ﬁrst tries to obtain optimal propensity score, and only then will it further tune the outcome model parameters to improve total loss (at the given optimal propensity score loss).
Thus for λ → ∞ optimizing (10) approximates the standard two-step approach that separates design and analysis stages of an experiment.
2.7. Binary Treatment & Gaussian Outcomes
In this section we describe parameter estimation for binary treatment and a continuous, univariate outcome.
3 It is also this scenario that we evaluate through simulation studies in Section 4.
Compute predicted j p j
and v i
j = ˆ p j
for j = 1 ,..., J : use weights v i , j to train a weighted outcome model for P ( Y | X , T ) .
The outcome models could be a neural net, a regression, or even just the sample mean of the response; Y | X , T = 1 and Y | X , T = 0 can also be modeled separately.
For binary treatment the negative log-likelihood for observation i is binary cross-entropy
[Equation] For univariate continuous outcomes a Normal log-likelihood, N ( y | x , σ 2 ε ) , is a natural candidate leading to a weighted sum of outcome losses for each state: [Equation]
[Equation]
where v i , j ( X i , θ T ) is the probability of unit i belonging to state j , ˆ Y
i ( θ Y ) is the outcome model prediction for Y i , and θ Y = { ( θ j , σ ε , j ) | j = 1 ,..., J } are all free parameters of the outcome model, with σ ε , j ∈ θ Y as the unknown standard deviation of the j -th outcome model residual, ε j : = Y i − ˆ Y i ( θ Y , j ) .
The average loss across N observations equals
[Equation]
It is worth to point out that the PRESS propensity model parameters, θ T , do not only enter (13) via the binary crossentropy propensity loss in (11), but also affect the outcome loss (12) through the weighting over the J states, v i , j ( X i , θ T ) .
Algorithm 1 describes an iterative learning procedure to optimize (13).
However, since PSPS models the joint distribution directly, it can be expressed as a single computational graph using standard neural network architectures (Fig. 2.1).
This has the beneﬁt of training all parameters and predictive states jointly, which avoids the – statistically and computationally – inefﬁcient iterative procedure.
In particular, the proposed MLE optimization and causal effect inference runs entirely within a single TensorFlow graph (Abadi et al., 2016), which in turn can be easily incorporated into larger graphs that rely on causal inference as part of their computation.
3. Estimating Causal Effects
Once the predictive state mapping and all outcome models have been trained, unit-level treatment effect (UTE) can be obtained as a weighted average of UTEs within each state:
[Equation]
where ˆ ∆ i ( j ) is the estimated UTE for unit i from outcome model j , and ε is discrete (continuous) for categorical (continuous) treatment.
Estimates for sample, population and conditional effects can be obtained by aggregating over the appropriate (sub-)sample of observations.
For example, average treatment effect (ATE) and average treatment effect on the treated (ATT) can be estimated as
[Equation]
ε
[Equation]
and [Equation]
3.1. Uncertainty Estimates
While in predictive learning tasks point predictions are of primary interest, in causal inference tasks getting accurate uncertainty estimates around the causal effect are equally – if not more – important.
A good causal estimator produces narrow conﬁdence intervals with proper coverage.
For simulations in Section 4, we use standard bootstrap procedures (Efron and Tibshirani, 1986) to obtain empirical conﬁdence intervals for ∆ i and ∆ ATE in (14) and (15), respectively.
4. Model Evaluation: Simulation Study
We compare PSPS to several state-of-the-art causal inference methods: an elastic net linear outcome model (Zou and Hastie, 2005), doubly robust inverse propensity weighted (DRIPW) models (Chan et al., 2010) – both linear and Random Forest (Breiman, 2001) versions –, entropy balancing (Hainmueller, 2012), a PRESS outcome model (Goerg, 2017) 4 and Dragonnet (Shi et al., 2019).
All of the above algorithms – with the exception of entropy balancing – can not only estimate the average treatment effect (ATE), but also unit-level treatment effects (UTE) (this is equivalent to the conditional average treatment effect (CATE) in our setup).
We conduct a simulation study on standard causal inference benchmarks and examine performance on ATE bias and UTE RMSE:
[Equation]
∑ n =( ∆
→ 1 i − ˆ ∆
→ 1 i ) 2 .
For each simulation we train a J = 10 PSPS model on n = 1 , 000 samples with an 80/20 train/test split and early stopping.
Figures 2 and 3 depict each run from a single simulated dataset as a point in the graph.
The boxplots give a sense of bias and spread of ATE, where a method with the same average (median) bias but smaller variance is preferred.
Similarly, the boxplots of RMSE summarize UTE estimation performance for each simulation.
This visualization is in lieu of the average bias as by construction the sample mean of the estimated UTEs for each dataset in the simulation is the estimated ATE for that dataset (see (15)); hence the average bias is equivalent to the bias for ATE. (a) ATE bias and UTE RMSE.
(b) 95% bootstrap conﬁdence interval coverage and width.
DRIPW linear model results are outliers in “Transformed” case and thus not shown for better visualization (CI width of 45 . 2; coverage of 0 . 7).
4.1. Kang & Schafer: Learning Non-Linear Causal Effects
Kang and Schafer (2007) present a simulation based on a real case study, where four features are drawn independently from a standard
Normal, z = ( z 1 , z 2 , z 3 , z 4 ) ∼ N ( 0 , I 4 ) , and the outcome and true propensity score π are generated as
[Equation]
) respectively ( ε ∼ N ( 0 , 1 ) ).
The original features ( z 1 , z 2 , z 3 , z 4 ) are then non-linearly transformed as x 1 = exp ( z 1 / 2 ) , x 2 = x 2 / ( 1 + exp ( z 1 )) + 10, x
3 = z 1 z 3 / 25 + 0 . 6 ) 3 , and x 4 = ( z 2 + z 4 + 20 ) 2 .
This allows us to evaluate causal methods against correctly and incorrectly speciﬁed models by either using the original ( z ) or transformed features ( x ).
This is particularly important for practitioners to see how well methods such as PSPS, PRESS, Dragonnet, and DRIPWRandomForest can approximate – in practice unknown – non-linear dependencies in real world data.
Figure 2 summarizes the results on bias & RMSE.
In the “Original” (linear) case, only entropy balancing and PSPS are unbiased; all other methods show medium to strong bias.
Interestingly, the elastic net LinearOutcomeModel is slightly biased on the causal estimate with very small RMSE, even though the true response surface is linear in the features.
We hypothesize that this can be attributed to incorrectly tuned (cid:96) 1 and (cid:96) 2 penalties in the elastic net cross validation, as it optimizes for predictive power, not accurate causal estimation.
In the “Transformed” case entropy balancing does signiﬁcantly worse on ATE bias.
6 PSPS on the contrary is effectively unbiased and has the smallest UTE RMSE amongst all methods.
Given that in practice, cause and effect are rarely – if ever – linearly related, this comes to show that
PSPS is a great choice for estimating causal effects from data collected in real world – non-linear – settings.
Figure 2 shows that only PSPS and entropy balancing have close to proper conﬁdence interval coverage with small widths in the “Original” case; in the “Transformed” case PSPS is the only algorithm that has proper conﬁdence intervals.
Interestingly, PSPS not only outperforms
[-I P W
-]
[Figure]
[Figure]
[Figure]
parametric & linear models but also signiﬁcantly improves on Dragonnet, which shows signiﬁcant bias for ATE and is unable to produce conﬁdence intervals at any reliability even in the linear case.
4.2. Lunceford & Davidian: Feature Selection
Lunceford and Davidian (2004) proposed another benchmark simulation where the data also consists of { ( X i , Z i , T i , Y i ) , i = 1 , ··· , n } , but all of them are observed.
Both X i and Z i are three dimensional vectors (see B.1 for details on data generating process for X and Z ).
The propensity score is only related to X – not Z – through π = expit (cid:0) β 0 + ∑ dj = 1 β j X i , j (cid:1) , where β controls the association between X and T .
We follow the original work and use β no = ( 0 , 0 , 0 , 0 ) T , β moderate =
( 0 , 0 . 3 , − 0 .
3 , 0 . 3 ) T , and β strong =
( 0 , 0 . 6 , − 0 .
6 , 0 . 6 ) T .
The response Y follows
Y i =
ν 0 + ∑ j =ν
j X ij
+ ν 4 T i + ∑ j =ξ j Z ij + ε i , ε i ∼ N ( 0 , 1 ) , ν = ( 0 , − 1 , 1 , − 1 , 2 ) T (17)
Similarly to β , ξ controls the association between Z and Y , ξ no = ( 0 , 0 , 0 ) T , ξ moderate = [-
-]( − 0 .
5 , 0 . 5 , 0 . 5 ) T , and ξ strong =
[-Conditional on X i 3 , ( X i 1 , Z i 1 , X i 2 , Z i 2 ) is then generated as multivariate normal N ( a X i 3 , B X i 3 ) , where a 1 = ( 1 , 1 , − 1 , − 1 ) T , a 
-]
( − 1 , 1 , 1 ) T .
The data generation process yields a true population ATE of ∆ ATE ≡ 2.
Estimating the ATE accurately is challenging in this simulation scenario given various inter- and intra-dependence of X and Z , while Z is only inﬂuencing the outcome – not the propensity score.
Figure 3 shows that PSPS, linear DRIPW, entropy balancing, and DRIPW perform equally well in all cases.
As for Kang & Schafer this is quite compelling for PSPS as it is on par with parametric linear models, while being able to learn non-linearities if data supports it.
Random Forest outcome models and Dragonnet increase in ATE bias for stronger associations.
Interestingly, both a linear and
Random Forest DRIPW model have difﬁculty estimating UTEs – compared to competing models with low RMSE.
Whilst consistently having the second lowest RMSE (after LinearOutcomeModel) the ATE bias of Dragonnet increases as the association between features and treatment increases.
Appendix B.2 has additional results from another causal simulation benchmark (Radcliffe, 2007). [Figure]
4.3. Balancing Outcome and Treatment Model Performance
A key aspect of PSPS is the joint optimization of the (treatment, outcome) likelihood.
The decomposition in (10) makes it clear that solely maximizing predictive power for the outcome will produce substantially worse causal estimates compared to an algorithm that strikes a balance between a good treatment model at the cost of some reduction in predictive power in the outcome models.
The simulation study allows us to replicate this empirically.
Figure 4 shows the R 2 of the underlying outcome model predictions ˆ y versus the ATE bias for each dataset.
As expected, methods with a higher R 2 are not guaranteed to achieve a better ATE bias reduction: PSPS often achieves zero bias even though it has a lower R 2 than other methods; vice-versa, even though RandomForestDRIPW and Dragonnet have the best predictive performance they fail to obtain unbiased ATE estimates.
4.4. Model Complexity and Parameter Tuning
PSPS builds on a J -state PRESS model to approximate P
( T | X ) , while simultaneously training outcome models in all J states at once.
The hyperparameter J controls the smoothness of the learned prediction functions.
The degenerate J = 1 case corresponds to a single outcome model with all observations falling in the same global state.
As J increases, the propensity score distribution can be approximated more closely while the sample size per state decreases at a rate of J − 1 .
Practitioners thus face a trade-off between better approximating P ( T | X ) by increasing J versus reducing J to avoid over-parameterization of the full model – especially the J outcome models ( θ 1 ,..., θ J ).
To explore how the choice of J affects PSPS performance we undertake a simulation study where we train PSPS models with J = 1 (degenerate), 2, 3, 5, 10, and 20.
Figure 5 illustrates that as long as J is in a reasonable range of O ( 1 ) − O ( 10 ) the causal estimation performance varies only slightly.
This aligns with the theoretical underpinning of PRESS as a kernel smoother that is quite robust to large J : as long as J is sufﬁciently large to obtain good coverage on the unit (propensity score) interval, the actual choice of J is secondary to obtain well-performing causal estimates. Simulations further showed that a more complex propensity model (e.g., several hidden layers feeding into predictive states) combined with simple (e.g., linear) outcome models is sufﬁcient to get unbiased causal effect estimates. We also found that shallow networks with standard activation functions help reduce variance of treatment effect estimates. While it remains a task for future work to explore whether PSPS is in fact a double-robust estimator (Robins and Rotnitzky, 1995), we
[Figure]
hypothesize that these ﬁndings can be explained by double-robustness as it is clearly preferable to get one propensity model right, rather than trying to specify all J outcome models correctly.
These results suggest a practical guideline on how to use PSPS: if ATE estimates do not change for J = 1 vs. J > 1, a simple (non-linear) outcome model might be sufﬁcient. This can be seen when comparing PSPS with J = 1 to other (non-PSPS) outcome model alternatives in each scenario. Interestingly, the coverage probability stays roughly constant for varying J , yet the conﬁdence interval width decreases signiﬁcantly as J increases. We hypothesize this is a result of lower variance estimates through the better outcome model approximation.
If practitioners are interested in obtaining narrow conﬁdence intervals, we recommend increasing the number of states.
Our ﬁndings are based entirely on the scenarios simulated here.
For real world datasets we recommend to run cross-validation and choose that J which maximizes the out-of-sample likelihood in (13).
5. Discussion
We introduce predictive state propensity subclassiﬁcation (PSPS), a novel causal inference algorithm that follows the Rubin Causal Model (RCM) framework by adjusting observational causal inference through propensity models following a subclassiﬁcation approach.
Rather than using adhoc procedures to stratify units by propensity score, PSPS estimates outcome models, propensity scores, and optimal strata simultaneously in a joint probabilistic model.
This yields efﬁcient training algorithms that can be easily implemented in deep learning frameworks like TensorFlow.
Several benchmark simulations show that PSPS has excellent statistical properties, with practically zero bias & low variance and the only method with proper conﬁdence interval coverage compared to several state-of-the-art causal inference algorithms in the statistical and machine learning literature. Acknowledgments
We want to thank Penny Chu, Tony Fagan, and Jim Koehler for their ongoing support during this project.
We are especially grateful to Mike Hankin, Tianqi Liu, Susanna Makela, Xiaojing Wang, and Mike Wurm for fruitful discussions and their contributions to run simulations and evaluations.
Appendix A. Methodology A.1. Deterministic Assignment to States
Without restricting w i , j ( X ) Algorithm 1 is not identical to traditional subclassiﬁcation, since each observation probabilistically maps to all states at once, rather than a deterministic x
i (cid:55)→ s j mapping.
In that sense PSPS is a soft-clustering generalization of classic subclassiﬁcation – akin to Gaussian mixture modeling (GMM) being a generalization of K-means (L¨ucke and Forster, 2019).
If practitioners want to obtain such a deterministic (hard-clustering) strata assignment, we suggest to add an entropy based penalty regularization of the form
[Equation]
N ∑
[Equation]
(cid:32) −
J ∑
[Equation]
j (cid:33)
[Equation]
N ∑
[Equation]
entropy i .
Adding (18) to (13) allows to control with how much certainty observations are mapped to predictive states.
For sufﬁciently large λ e (cid:29)
(18) is minimized for deterministic mappings, w i , j = 0 for all but one j for each i , i.e., it assigns observations to exactly one state with probability 1.
We want to emphasize though that without any prior knowledge of a deterministic mapping we suggest to keep the soft-thresholding property of PRESS, as from a likelihood point of view it leads to better models.
Similarly to how a GMM usually gives higher quality results compared to a hard-thresholding K-means.
In practice a small λ e can be beneﬁcial to make w j more sparse.
A.2. Trimming
Using predictive states in the propensity model allows us to easily assess “treatment vs. control” balance within each predictive state and then drop states from inference for which balance is deemed inadequate.
A simple suggestion is to ensure that there are enough treated and control units in each state with respect to the complexity of the outcome model; alternatively balance of individual features or their products may be assessed via methods such as t-tests, etc.
Appendix B. Simulation Studies B.1. Feature generation for Lunceford & Davidian
{+The joint distribution of ( X i , Z i ) is given by X i 3 ∼ Bernoulli ( 0 . 2 ) , then generate Z i 3 as Bernoulli with P ( Z i 3 = 
 | X i 3 )+}
= 0 .
75 · X i 3 + 0 . 25 · ( 1 − X i 3 ) .
{+Conditional on X i 3 , ( X i 1 , Z i 1 , X i 2 , Z i 2 ) is then generated as multivariate normal N ( a X i 3 , B X i 3 ) , where a 1 = ( 1 , 1 , − 1 , − 1 ) T , a 
+}
= ( − 1 , − 1 , 1 , 1 ) T and
[Equation]
B.2. Additional Simulation Study: Radcliffe & Surrey
This data generating process is taken from Radcliffe (2007) and Radcliffe and Surry (2011).
Here treatment is randomly assigned to each unit with probability p .
For each unit, x 1 and x 2 are drawn
[Figure]
[Figure] [Figure] ) as independent uniform random integers in U [ 0 , B ] for an upper bound B (here: B = 99).
Unobserved variables u 1 and u 2 are then drawn uniformly from U ( 0 , x i ) , i = 1 , 2 respectively.
Unit-level treatment effects and outcome regressions are then set as
[Equation]
where T i is the realization of a Bernoulli draw with probability p (here: p = 0 . 5).
This speciﬁcation yields a true ATE of
ATE = s · E ( u 2 ) + c , where E ( u 2 )
=
B / 4 .
In simulations we explore whether the scaling factor s or shift c affect the ability to recover the true treatment effect.
As for Lunceford & Davidian we consider three scenarios: const = { s = 0 , c = 3 } , scale = { s = 0 .
1 , c = 0 } , and both = { s = 0 .
1 , c = 3 } .
Figure 6 shows that – except for the outcome PRESS neural net – all models give unbiased ATE estimates; and again, linear and Random Forest DRIPW have difﬁcult estimating unit-level effects.
Results do not vary much between “const”, “scale”, or “both” except for RandomForestDRIPW, which performs worse as the association grows stronger.
B.3. Propensity Metrics and Near Positivity Violations
A common concern with propensity score methods, in particular inverse propensity weighting (IPW), are near positivity violations which are situations where the propensity to treat is close toor 1 (Robins et al., 2007).
Practically this means it is unlikely to ﬁnd any matching control (treated) units.
Statistically IPW algorithms run into near zero-division issues which result in high instability in the causal effect estimate.
We explore how estimators perform across two metrics on the propensity score distributions that measure important properties related to near-positivity violations:
• Extremeness : proportion of scores outside of [ 0 . 05 , 0 .
This highlights scenarios where the propensity model is close to near-positivity violations (Robins et al., 2007).
• Flatness : normalized Shannon entropy of the propensity score distribution (unconditional) (Cover and Thomas, 2006).
This measures how ﬂat (or peaked) the propensity distribution is.
A score of 1 . 0 is a uniform distribution; a lower score corresponds to a more “peaky” distribu
- tion.
Exactly zero score ﬂatness occurs when the propensity scores reduce to singletons along the unit interval, and thus the continuous distribution degenerates to delta functions.
This can happen for PSPS if the mapping from features to states ( p ( s j | X i ) ) become deterministic, and only J propensity scores in the ﬁnal propensity model remain.
Across all simulations we only found that one scenario resulted in non-trivial relationships between the propensity score distribution metrics above and the quality of the ATE estimates.
For the original KangSchafer dataset, without transformations, we see that DoublyRobustIPW shows various degrees of ﬂatness and extremeness, which result in improved ATE estimates (as ﬂatness and extremeness increased).
For PSPS on the other hand all scores are away from near-one violations (extremeness = 0 .
for all simulations) and propensity score ﬂatness is close to 0 .
– which corresponds to predictive state mappings that are deterministic functions from features, X , to one of the J states, s 1 ,..., s J .
In either case, it had no impact on quality of ATE estimate in terms of bias.
An explanation for this lies in how PSPS uses (estimated) propensity scores to arrive at the ﬁnal causal estimate: it computes a convex-combination, weighted average of the J individual outcomes models.
The key here is that propensity scores are not inverted, as in any inverse-propensity weighting method, thus near-one (or near zero) propensity scores pose no problem for PSPS. [Figure] [Figure]
B.4. Sensitivity Analysis For Number of States
In most scenarios J = 10 was sufﬁcient, and any further increases in J only yielded marginal (if any) improvements.
In fact, Figure 5 shows that increasing J decreases RMSE while in Fig.
8 decreasing J is beneﬁcial with J = 10 a natural midpoint between the two.
Even a single state PSPS architecture ( J = 1 ) is competitive in certain scenarios.
This is an interesting edge case as it is equivalent to a simple outcome model with averaging weights w i , j = w i , 1 ≡ 1 for each i .
[-D R I P W2 48 1014 
-]
[-D R I P W2 48 1014 
-]
[-D R I P W2 48 1014 
-]
[Figure]
[Figure]
[Figure]