Propensity Score Stratification in Observational Data: A Tutorial | by Leihua Ye, PhD | Towards Data Science

Do you know how much new data we generate each day? In 2021, the number is roughly 2.5 quintillion bytes of data. So, we are not in short supply of data but lack the correct tool of making it useful.
It’s easy to show how multiple factors move together with access to big data. However, it becomes more challenging to determine whether the relationship is causal or merely correlational. At large, there are three types of designs that can help us:
True Experimental Designs, aka A/B tests, or Randomized Controlled TrialsQuasi-Experimental DesignsObservational Designs
True Experimental Designs, aka A/B tests, or Randomized Controlled Trials
Quasi-Experimental Designs
Observational Designs
In a series of blog posts, I’ve elaborated on the related topics within each category. For example, we have to understand the value of A/A tests, avoid common pitfalls, follow best practices, and be careful with user interference while running A/B tests.
There are various available designs in the quasi-experimentation category, including Difference-In-Difference, Regression Discontinuity Design, Interrupted Time Series, Synthetic Control, and Instrumental Variable. A common misleading argument is X/Y/Z method is the best design. Instead, data scientists have to assess the available information and choose the most appropriate one to the circumstance.
In the last category, Propensity Score holds a central role in obtaining causal inference using observational designs (Rosenbaum, 1983). Propensity Score Analysis has four main methods: PS Matching, PS Stratification, PS Weighting, and Covariate Adjustment.
In a prior post, I’ve introduced how we can use PS Matching to reduce the observed baseline covariate imbalance between the treatment and control groups. In today’s post, we elaborate on how to stratify observations on the basis of PS and explain why it is a better alternative to PS Matching. In the empirical part, we code it up in R and calculate the treatment effects using a real-life example.
PS is a scaled score of the conditional probability of receiving the treatment given the observed covariates. Thus, P.S. is a balancing metric that effectively checks covariate balance between the treatment and control groups.
A potential risk to any PS Analysis is the presence of unobserved covariates that affects the outcome variable. There are three ways of testing if the unobserved covariates are present (Rosenbaum, 2005):
Construct two control groups that are known to differ on the unobserved covariates.Check coherence of associations and dose-response relationships.Check the variables that should not be affected by the treatment.
Construct two control groups that are known to differ on the unobserved covariates.
Check coherence of associations and dose-response relationships.
Check the variables that should not be affected by the treatment.
Please check Rosenbaum (2005) for more information.
Let’s address the elephant in the room:
What’s the fuss about Propensity Score?
Why does every single observational method have to play with it?
First off, we have to understand the nature of observational data and why it is difficult to do causal inference using non-experimental (observational) data. The reason why we call it observational data is because data scientists have no control over or knowledge of the treatment assignment process. We merely observe some units following into the treatment group and others in the control group.
The absence of valid intervention makes it difficult to establish creditable causal relationships. In RCTs, we adopt the difference-in-two-means estimator to quantify the intervention effect, i.e., comparing the mean differences between the treatment and control groups. However, we can’t use the same estimator in observational data, which may contain a selection bias. That is, people self-select into the treatment group and others into the control group.
Here is an example.
Economists want to understand the effectiveness of job training programs in affecting participants’ household income. However, if we simply compare the post-program salaries between program participants and non-participants, we would end up with a biased estimate because those people who are more motivated than the average are more likely to enroll in the training program.
If we merely compare the differences without adjustment, we may attribute the effects caused by the cofounding variable (better motivation, in this case) to the intervention effect:
To obtain an unbiased estimator, we have to adjust for the key covariate imbalance and reduce (or even eliminate) the selection bias.
In the next section, let’s dig into the reasons why PS Stratification.
PS is a balancing score: conditional on the PS, the distribution of the observed covariates looks similar between treated and control groups (Austin, 2011). Thus, it allows you to adjust for the covariate imbalance by tweaking the score.
Some researchers argue that we can match participants based on their PS and find comparable cases, i.e., PS Matching. However, the precise matching process increases imbalance, inefficiency, model dependence, bias and fails to reduce the imbalance (King and Nielsen, 2019). In contrast, PS Stratification offers a better alternative to PS Matching.
Here are the specific steps:
1. Estimate the PS using a logistic regression 2. Create mutually exclusive strata based on the estimated PS3. Group treated and control units into each stratum 4. Within each stratum, we calculate the difference in means between the treated and control 5. Weight the means in each stratum and obtain the target estimate
There are two additional caveats with step 2 and step 5. At step 2, research shows 5 strata can remove approximately 90% of the bias in the unadjusted estimate (Rosenbaum and Rubin, 1984). A reasonable assumption is: What if we increase the strata as many as possible? Would it lead to a better bias reduction?
No.
Simulation works have shown that 5–10 strata lead to the best result, and 10+ strata bringing little marginal gain (Neuhäuser et al., 2018). Also, there is a practical reason: the more strata we have, the fewer data points will lie within each stratum.
At step 5, PS stratification allows us to estimate both Average Treatment Effect (ATE) and Average Treatment Effect on the Treated (ATT) depends on how we weigh the means.
To estimate the ATE, we weigh by the number of units in each stratum;
To estimate the ATT, we weigh by the number of treated units in each stratum (Imbens, 2004).
You may be confused with the Potential Outcomes Framework (POF) notations: what do ATE and ATT actually mean?
ATE: the treatment effect averaged over the entire population.
ATT: comparing the treated subjects to similar subjects who are not treated.
Conceptually, PS stratification is considered a meta-analysis of a set of quasi-RCTs within each stratum (Austin, 2011). We can think of the treatment assignment process as a random assignment conditional on observed covariates, aka the ignorable assumption, meaning the potential outcomes are independent of the treatment status (Austin, 2011).
Within the POF, the ignorable assumption allows us to obtain an unbiased estimate.
where
Y0i: if unit i does not receive the treatment
Y1i: if unit i receives the treatment
Zi: treatment condition
Usually, the ignorable assumption does not hold in observational studies. In other words, units in the treated groups are more likely to receive the treatment than those in the control. However, PS Stratification offers a solution: conditional on the observed covariates represented as the PS, the ignorable assumption holds, i.e., Conditional Ignobility Assumption (CIA).
The two-group-means-difference estimator is the unbiased stratum-specific estimate of the ATE (D’Agostino, 1998).
where
Y0i: if unit i does not receive the treatment
Y1i: if unit i receives the treatment
Zi: treatment condition
Xi: a vector of observed covariates represented as the PS
That’s all the theories behind the PS Stratification.
Let’s code it up in R.
In this section, we use a composite dataset used by LaLonde (1986) and Dehejia and Wahba (1997). It constitutes of two parts: a subsample of the treated data in the National Supported Work Demonstration (NSW) and a comparison sample from the Population Survey of Income Dynamics (PSID).
For pedagogical reasons, I simplify the analysis process and mainly focus on the workflow of how to conduct a Propensity Score Stratification in R. That’s why the obtained ATT is slightly different from the original papers.
A quick recap of the PS Stratification:
1. Estimate the PS using a logistic regression 2. Create mutually exclusive strata based on the estimated PS3. Group treated and control units into each stratum 4. Within each stratum, we calculate the difference in means between the treated and control 5. Weight the means in each stratum and obtain the target estimate
Step 0: Data preparations
#install.packages(“MatchIt”)#install.packages(“optmatch”)library(“MatchIt”)library(“optmatch”)data(“lalonde”)head(lalonde)#?lalonde
```{r}sum(lalonde$treat)```185
The number of total treated cases is 185.
```{r}mean(lalonde[which(lalonde$treat==0),’re78'])```6984.17
The outcome variable in the control group is $6984.17.
```{r}mean(lalonde[which(lalonde$treat==1),’re78'])```6349.144
The outcome variable in the treated group is $6349.144.
Step 1: Estimate the PS using a logistic regression
```{r}PS_logit <- glm(treat ~ age + educ+race+married+nodegree+re74+re75, data = lalonde, family=binomial)PS_Scores = fitted(PS_logit)head(PS_Scores)```
Step 2: Create mutually exclusive strata based on the estimated PS
```{r}lalonde$PS_Scores <- PS_Scoreslalonde```
```{r}# create strata Quintiles <- quantile(lalonde$PS_Scores, prob=seq(from=0,to=1,by=0.2),na.rm=TRUE)lalonde$PS_Stratum <- cut(lalonde$PS_Scores, breaks = Quintiles, labels = 1:5, include.lowest = TRUE)lalonde```
Step 3 & 4: Group treated and control units into each stratum; Within each stratum, we calculate the difference in means between the treated and control
In R, we can combine the steps 3 and 4 together by running the following code:
```{r}# Fifth Stratum: PS_Stratum = 5diff_5 = mean(lalonde[which(lalonde$treat==1&lalonde$PS_Stratum==5), ’re78'])- mean(lalonde[which(lalonde$treat==0&lalonde$PS_Stratum==5),’re78'])diff_5```1364.036
In Stratum 5, the difference of the outcome variable (‘re78’) between the treated and the control units is $1364.036.
We repeat the same process for the remaining strata in the following code:
```{r}# Forth Stratum: PS_Stratum = 5diff_4 = mean(lalonde[which(lalonde$treat==1 & lalonde$PS_Stratum==4),’re78']) — mean(lalonde[which(lalonde$treat==0 & lalonde$PS_Stratum==4),’re78'])diff_4841.7422# Third Stratum: PS_Stratum = 3diff_3 = mean(lalonde[which(lalonde$treat==1 & lalonde$PS_Stratum==3),’re78']) — mean(lalonde[which(lalonde$treat==0 & lalonde$PS_Stratum==3),’re78'])diff_33167.41# Second Stratum: PS_Stratum = 2diff_2 = mean(lalonde[which(lalonde$treat==1 & lalonde$PS_Stratum==2),’re78']) — mean(lalonde[which(lalonde$treat==0 & lalonde$PS_Stratum==2),’re78'])diff_22122.768# First Stratum: PS_Stratum = 1diff_1 = mean(lalonde[which(lalonde$treat==1 & lalonde$PS_Stratum==1),’re78']) — mean(lalonde[which(lalonde$treat==0 & lalonde$PS_Stratum==1),’re78'])diff_1-10467.06```
As it should be pointed out, the difference in stratum 1 appears to be off: -10467.06. The reason is there is only one treated observation in the stratum, and technically we should leave it out. However, for pedagogical reasons, we keep it in the analysis.Step 5. Weight the means in each stratum and obtain the target estimate, ATT
```{r}# the total nunmber of treated cases total_treated_n <- nrow(lalonde[which(lalonde$treat==1),])total_treated_n 185```
First off, we calculate the total number of treated cases, and there are 185.
```{r}#the number of treated cases in each stratum stratum_5_treated_n <- nrow(lalonde[which(lalonde$treat==1 & lalonde$PS_Stratum==5),])stratum_5_treated_n 85```
In Stratum, there are 85 treated cases.
```{r}ATT_5 = diff_5*stratum_5_treated_n/total_treated_nATT_5626.7194```
Then, we calculate the weighted treatment effect in Stratum 5.
Repeat the same procedure for Statum 1–4.
```{r}#the number of treated cases in each stratum # stratum 4stratum_4_treated_n <- nrow(lalonde[which(lalonde$treat==1 & lalonde$PS_Stratum==4),])stratum_4_treated_nATT_4 = diff_4*stratum_4_treated_n/total_treated_nATT_4323.047# stratum 3stratum_3_treated_n <- nrow(lalonde[which(lalonde$treat==1 & lalonde$PS_Stratum==3),])stratum_3_treated_nATT_3=diff_3*stratum_3_treated_n/total_treated_nATT_3359.5438# stratum 2stratum_2_treated_n <- nrow(lalonde[which(lalonde$treat==1 & lalonde$PS_Stratum==2),])stratum_2_treated_n7# stratum 1stratum_1_treated_n <- nrow(lalonde[which(lalonde$treat==1 & lalonde$PS_Stratum==1),])stratum_1_treated_nATT_1= diff_1*stratum_1_treated_n/total_treated_nATT_1- 56.57873```
As a last step, we add it up and obtain the ATT:
```{r}ATT = ATT_1+ATT_2+ATT_3+ATT_4+ATT_5ATT1333.052```
Using PS Stratification, we find that attending the job training program leads to a $1333.052 increase in after-program household income. The experimental benchmark is $1,794, and the original paper estimates 1,608 using the PS Stratification. The gap is due to some technical details.
Medium recently evolved its Writer Partner Program, which supports ordinary writers like myself. If you are not a subscriber yet and sign up via the following link, I’ll receive a portion of the membership fees.
leihua-ye.medium.com
Austin, P.C., 2011. An introduction to propensity score methods for reducing the effects of confounding in observational studies. Multivariate behavioral research, 46(3), pp.399–424.
D’Agostino Jr, R.B., 1998. Propensity score methods for bias reduction in the comparison of a treatment to a non‐randomized control group. Statistics in medicine, 17(19), pp.2265–2281.
Dehejia, R.H. and Wahba, S., 1999. Causal effects in non-experimental studies: Reevaluating the evaluation of training programs. Journal of the American statistical Association, 94(448), pp.1053–1062.
Imbens, G.W., 2004. Nonparametric estimation of average treatment effects under exogeneity: A review. Review of Economics and statistics, 86(1), pp.4–29.
King, G. and Nielsen, R., 2019. Why propensity scores should not be used for matching. Political Analysis, 27(4), pp.435–454.
LaLonde, R.J., 1986. Evaluating the econometric evaluations of training programs with experimental data. The American economic review, pp.604–620.
Neuhäuser, M., Thielmann, M. and Ruxton, G.D., 2018. The number of strata in propensity score stratification for a binary outcome. Archives of medical science: AMS, 14(3), p.695.
Rosenbaum, P.R. and Rubin, D.B., 1984. Reducing bias in observational studies using subclassification on the propensity score. Journal of the American statistical Association, 79(387), pp.516–524.
Rosenbaum, P.R., 2005. Sensitivity analysis in observational studies. Encyclopedia of statistics in behavioral science.
Please find me on LinkedIn and Youtube.
Also, check my other posts on Artificial Intelligence and Machine Learning.