📊 statistics

Bayesian Nonparametrics for Principal Stratification with Continuous Post-Treatment Variables

This paper introduces CASBAH, a novel Bayesian nonparametric method that utilizes a hierarchical structure to identify data-adaptive coarsened principal strata and quantify uncertainty for causal inference with binary treatments and continuous post-treatment variables, demonstrating superior performance in simulations and a real-world application to US air quality regulations.

Original authors: Dafne Zorzetto, Antonio Canale, Fabrizia Mealli, Francesca Dominici, Falco J. Bargagli-Stoffi

Published 2026-03-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Dafne Zorzetto, Antonio Canale, Fabrizia Mealli, Francesca Dominici, Falco J. Bargagli-Stoffi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a doctor trying to figure out if a new medicine actually works. You give it to some patients (the "treatment" group) and not others (the "control" group). But here's the catch: the medicine doesn't just affect the final health outcome; it first changes something in the middle, like a blood test result.

In statistics, this is called Principal Stratification. It's like trying to sort patients into different "teams" based on how their bodies would have reacted to the medicine, regardless of whether they actually got it.

Team A (The Responders): Their blood test would change whether they got the medicine or not.
Team B (The Non-Responders): Their blood test would stay exactly the same, no matter what.

The problem? You can't see the "what if." You only see what happened. If a patient didn't get the medicine, you don't know what their blood test would have been if they had.

The Old Way: The "Rough Cut" Problem

For years, statisticians tried to solve this with continuous data (like blood pressure or pollution levels, which are numbers on a scale, not just "high" or "low") by using a Rough Cut.

Imagine you have a giant bucket of water (the data). To sort it, you decide: "Anything above 50 degrees is 'Hot,' and anything below is 'Cold'."

The Flaw: This is arbitrary. Why 50? Why not 49? If you pick the wrong number, you might put a "Hot" person in the "Cold" bucket, ruining your analysis. It's like trying to sort a rainbow by cutting it with a pair of scissors; you lose the nuance.

The New Way: CASBAH (The "Smart Sorter")

The paper introduces a new method called CASBAH (Confounders-Aware SHared atoms BAyesian mixture). Think of CASBAH as a super-smart, shape-shifting sorter that doesn't need scissors.

Here is how it works, using a few analogies:

1. The "Shared Atoms" (The Lego Bricks)

Imagine the data is made of invisible Lego bricks.

In the old methods, the "Hot" group and the "Cold" group were built with completely different sets of bricks.
CASBAH says: "Let's use the same set of Lego bricks to build both groups."
Because the groups share the same building blocks, the model can naturally figure out which people belong to which group based on how the bricks fit together. It doesn't need a ruler to cut them; it just sees which shape they naturally form.

2. The "Data-Adaptive" (The Clay Sculptor)

Instead of forcing the data into a square or a circle (like the old "Rough Cut" method), CASBAH is like a sculptor working with wet clay.

It looks at the data and says, "Oh, this group of people naturally forms a long, thin shape. And this other group forms a round blob."
It molds the "teams" (strata) based on the actual shape of the data, not a rule you made up beforehand. This means it can find the "Non-Responders" (the dissociative stratum) even if they are hidden in the middle of the data, without needing to guess a cutoff point.

3. The "Confounder Awareness" (The Detective's Notebook)

Sometimes, a person's age, income, or location (confounders) affects how they react.

CASBAH keeps a detective's notebook. It looks at all these extra details for every single person and uses them to decide which "team" they are most likely on. It's not just looking at the blood test; it's looking at the whole picture to make a fair guess.

The Real-World Test: Cleaning the Air

To prove it works, the authors applied CASBAH to a real-life problem: US Air Quality Regulations.

The Setup: Some counties were told to clean up their air (Treatment), and others weren't (Control).
The Middle Variable: The actual level of pollution (PM2.5).
The Outcome: Mortality rates (deaths).

What they found:
Using the old "Rough Cut" methods, it was hard to tell who actually benefited. But CASBAH sorted the counties into three clear teams:

The "Cleaners" (Associative Negative): These counties actually reduced their pollution because of the rules. Result: Their death rates went down. The medicine worked!
The "Unchanged" (Dissociative): These counties were supposed to clean up, but their pollution levels didn't change at all. Result: Their death rates didn't change either. The rules didn't work here.
The "Worseners" (Associative Positive): A few places actually got more polluted (maybe due to weird weather or other factors). Result: No improvement in health.

Why This Matters

The big takeaway is flexibility and honesty.

No Arbitrary Lines: You don't have to guess where to draw the line between "good" and "bad" pollution. The math finds the natural groups.
Uncertainty is Okay: CASBAH admits, "I'm 80% sure this county is a 'Cleaner' and 20% sure it's an 'Unchanged'." It gives you a probability, not just a yes/no answer. This helps policymakers understand that not every situation is black and white.

In short, CASBAH is a new statistical tool that lets us sort people into meaningful groups based on how they would react to a treatment, even when the data is messy and continuous, without forcing it into a box that doesn't fit. It's like upgrading from a pair of scissors to a sculptor's hands.

1. Problem Statement

Principal stratification is a causal inference framework used to estimate treatment effects in the presence of post-treatment (intermediate) variables. While well-established for binary post-treatment variables (yielding a finite number of strata, typically four), the framework faces significant challenges when the post-treatment variable is continuous.

The Challenge: With a continuous intermediate variable, the set of principal strata becomes potentially infinite, indexed by the joint potential responses $(P_i(0), P_i(1))$ .
Limitations of Existing Methods:
- Dichotomization: Converting continuous variables to binary via arbitrary thresholds (e.g., $P_i(1) > P_i(0)$ ) loses information and introduces subjectivity.
- Parametric/Ad-hoc Approaches: Existing methods often rely on restrictive parametric assumptions or externally imposed clustering rules to coarsen the infinite strata.
- The "Dissociative" Stratum: A critical methodological gap exists in defining the "dissociative stratum" (units where the treatment has no effect on the intermediate variable, i.e., $P_i(1) = P_i(0)$ ). In continuous settings with diffuse distributions, the probability of exact equality is zero, making this stratum unidentifiable without arbitrary thresholds.

2. Methodology: CASBAH Model

The authors propose CASBAH (Confounders-Aware SHared-atoms BAyesian mixture), a novel Bayesian nonparametric approach designed to handle continuous post-treatment variables.

Core Architecture

Dependent Dirichlet Process (DDP): The model utilizes a DDP prior (specifically a probit stick-breaking process) to model the distribution of potential post-treatment outcomes $P_i(t)$ conditional on covariates $X_i$ and treatment $t \in \{0, 1\}$ .
Shared Atoms: A key innovation is the use of shared atoms across treatment levels. The random probability measures $G^{(0)}_{x_i}$ $G_{x_{i}}^{(0)}$ and $G^{(1)}_{x_i}$ $G_{x_{i}}^{(1)}$ share the same sequence of kernel parameters $\{\psi_l\}$ ${ψ_{l}}$ .
- $P_i(t) | x_i, t \sim \sum_{l} \pi^{(t)}_l(x_i) K(\cdot; \psi_l)$
- The weights $\pi^{(t)}_l(x_i)$ are treatment-specific and depend on covariates via a probit regression, but the "atoms" (cluster centers) are shared.
Latent Strata Definition: Principal strata are defined by the latent cluster allocation variables $S^{(0)}_i$ $S_{i}^{(0)}$ and $S^{(1)}_i$ $S_{i}^{(1)}$ .
- Dissociative Stratum ( $S_0$ ): Occurs when $S^{(0)}_i = S^{(1)}_i$ (the unit belongs to the same cluster under both treatment and control).
- Associative Strata ( $S_+, S_-$ ): Occur when $S^{(0)}_i \neq S^{(1)}_i$ , indicating a shift in the distribution of the intermediate variable.

Key Theoretical Properties

Non-Zero Prior Probability: Unlike continuous models where $P(P_i(1)=P_i(0))=0$ , the discrete nature of the Dirichlet process support ensures that $P(S^{(0)}_i = S^{(1)}_i) > 0$ a priori. Theorem 1 provides a closed-form expression for this probability, showing it depends on covariates $x$ .
Confounder Adjustment: The mixture weights are modeled as functions of observed covariates, allowing for flexible, nonparametric adjustment for confounding.
Uncertainty Quantification: The framework jointly imputes missing potential outcomes and propagates uncertainty regarding stratum membership through the posterior distribution.

Inference

Algorithm: Posterior inference is conducted via Gibbs sampling.
Steps:
1. Sample latent cluster allocations $S^{(t)}_i$ .
2. Sample cluster-specific parameters (location and scale of the Gaussian kernels).
3. Sample regression parameters for the mixture weights.
4. Impute missing potential post-treatment variables ( $P_i(1-T_i)$ ).
5. Sample parameters for the primary outcome model (assumed linear in the paper, but extensible).

3. Key Contributions

Data-Adaptive Strata Identification: Eliminates the need for arbitrary thresholds or pre-specified discretization. Strata are identified endogenously based on the data structure.
Probabilistic Dissociative Stratum: Solves the identification problem for the dissociative stratum in continuous settings by leveraging the discrete support of the Dirichlet process, assigning a non-zero probability to units where the treatment does not alter the intermediate variable.
Comprehensive Uncertainty: Provides full posterior distributions for principal causal effects (PCE), accounting for uncertainty in stratum membership and missing data imputation.
Covariate-Dependent Mixing: Integrates confounders directly into the mixture weights, offering a flexible alternative to rigid parametric models.

4. Results

Simulation Study

The authors compared CASBAH against Schwartz et al. (2011) (parametric) and Lu et al. (2023) (copula-based) across five scenarios with varying strata configurations (dissociative, positive associative, negative associative) and numbers of confounders.

Bias: CASBAH demonstrated significantly lower bias in estimating the expected difference in post-treatment variables ( $E[P(1)-P(0)]$ ) and primary outcomes ( $E[Y(1)-Y(0)]$ ) compared to competitors.
Strata Identification: The Adjusted Rand Index (ARI) for CASBAH was consistently high (0.91–0.99), indicating accurate recovery of the true principal strata. Competing models failed to identify strata according to the defined causal framework.
Robustness: CASBAH performed well even when strata atoms were close together or variances differed.

Case Study: US Air Quality Regulations (NAAQS)

Context: Estimating the causal effect of the 2005 NAAQS revision (treatment) on mortality rates (outcome), mediated by $PM_{2.5}$ levels (continuous post-treatment variable).
Data: 384 counties in the Eastern US (270 attainment, 114 non-attainment).
Findings:
- Strata Identification: The model identified three distinct groups:
  - Dissociative (32%): Pollution levels unchanged by regulation.
  - Associative Negative (56%): Pollution significantly reduced.
  - Associative Positive (12%): Pollution unexpectedly increased.
- Causal Effects:
  - In the Associative Negative stratum (where pollution dropped), the mortality rate decreased significantly (median reduction of 8.12 per 10,000).
  - In the Dissociative and Associative Positive strata, no significant reduction in mortality was observed.
- Characterization: The model revealed that counties in the "Associative Positive" stratum were often rural, had lower income, and higher smoking rates, whereas "Associative Negative" counties were urban with higher education and income.

5. Significance and Conclusion

This paper represents a significant advancement in causal inference for continuous intermediate variables. By moving away from arbitrary discretization and rigid parametric assumptions, CASBAH offers a principled, flexible, and interpretable framework for principal stratification.

Theoretical Impact: It resolves the theoretical impossibility of defining a dissociative stratum in continuous settings without ad-hoc rules.
Practical Impact: The application to environmental policy demonstrates that treatment effects are heterogeneous and dependent on the mechanism of change (i.e., whether the policy actually reduced pollution). This allows policymakers to distinguish between areas where regulations are effective versus those where they are not, preventing misleading aggregate conclusions.
Future Directions: The authors suggest extending the outcome model to non-linear structures and adapting the framework for continuous treatments or time-to-event settings.