A subtle selection bias in machine learning-based causal inference can flip your results — here’s how to spot and fix it
Machine learning is great at predicting outcomes. But when we try to use it to estimate causal impact, especially in a two-stage setup, things can go surprisingly wrong – even when the treatment was randomly assigned.
In this article, I’ll guide you through a subtle but critical pitfall that often arises when using a two-stage machine learning model to estimate causal uplift, particularly in marketing and product experimentation settings. This isn’t just a theoretical concept – it’s an easy mistake to make, and it can completely reverse the apparent direction of your treatment effect.
I’ll explain:
- Why a two-stage ML model is tempting (and sometimes necessary)
- A real-world example where it fails
- The underlying selection bias that breaks the logic
- How to fix it and get back to an unbiased estimate
Why Use Machine Learning to Estimate Causal Impact?
Let’s say you’re running an A/B test on an ad campaign. Treatment is randomly assigned, so the true causal impact should be:
E[Y | T = 1] − E[Y | T = 0]
Where Y is your outcome – say, the dollar amount a user spends.
In practice, though, especially in digital marketing or e-commerce, most users don’t make a purchase. So for the majority, Y = 0. This introduces a lot of variance and obscures the real signal you care about: how much the ad is influencing actual buyers.
To deal with this, many teams opt for a two-stage model setup:
- Stage 1: Predict whether the user will make a purchase (Y > 0)
- Stage 2: Predict the amount of purchase (Y) only for those users predicted to buy
This is especially appealing when actual buyers are rare, and you want a more focused estimate of incremental revenue, but only for those who make a purchase.
A Real-World Example: When the Model Backfires
Imagine you follow the two-stage setup: First, you train a model to predict the likelihood of each user to make a purchase. Then, you select users whose predicted probability is greater than 50%. You apply a second model to estimate revenue (Y) for these selected users and compare the test group (treatment) with the control group:
E[Y | T = 1, Y > 0] − E[Y | T = 0, Y > 0]
You might be surprised to see a negative uplift, as though the treatment actually hurts sales! But you know that’s impossible—the treatment was randomly assigned. So, what went wrong?
The Hidden Bias: Conditioning on an ML Filter
Here’s what’s happening behind the scenes:
Since the treatment group saw an ad, some marginal users were nudged into buying—maybe they bought a cheap trial pack or sample-size product. These light buyers are included in the treatment group because the Stage 1 model predicted they’d make a purchase.
But in the control group, these marginal users aren’t influenced by the ad so stage 1 model predicts no purchase.
As a result:
- Treatment group contains many low-spending, new buyers (who were nudged by the ad).
- Control group mostly contains high-likelihood, high-spending buyers (the ones that model predicted would buy anyway)
This skews the comparison. The average revenue in the treatment group is lower, even though the treatment led to more purchases!
You just fell into a selection bias trap, created by your model.
The Right Way to Use Two-Stage Models for Causal Estimation So, how do you get valid causal estimates from a two-stage model?
The key is to avoid conditioning on post-treatment variables (like in this case model based outcome Y_hat > 0) which varies between treatment and control.
Instead:
- Train both models as before.
- At inference time, compute expected sales for each user, regardless of predicted purchase.
Ŷfinal = P( Y>0 ) × E[ Y ∣ Y>0 ]
Here:
- P( Y>0 ) comes from Stage 1 (classification).
- E[ Y ∣ Y>0 ] comes from Stage 2 (regression).
You multiply the two values, which gives you the unconditional expected value of Y. This ensures your model respects the original randomized design and avoids introducing bias into the comparison.
Then estimate causal impact the right way:
E[ Ŷfinal ∣ T=1 ] − E[Ŷfinal ∣ T=0]
This gives you the true causal estimate without the bias from your model.
Conclusion
It’s easy to get caught up in the allure of sophisticated models that seem to “clean up” your data before analysis. But when it comes to causal inference—especially with randomized treatment—you have to be extra careful about introducing post-treatment bias.
Your model’s job is to support your estimate — not decide who gets included in the analysis.
So the next time you build a two-stage ML model to estimate causal impact, don’t just check if your predictions are good. Ask yourself: Does your pipeline preserve the integrity of your randomized experiment?
Because, when it comes to measuring impact, how you slice the data is often more important than the model itself.
Author: Palash Arora