Add Sensitivity bounds for ATE via Marginal Sensitivity Model#925
Add Sensitivity bounds for ATE via Marginal Sensitivity Model#925aman-coder03 wants to merge 1 commit into
Conversation
jeongyoonlee
left a comment
There was a problem hiding this comment.
Two blockers before merge, both in how potential outcomes are pulled from the learner:
B1 — get_potential_outcome_predictions misreads X-learner components (and the test uses the X-learner).
The method assumes fit_predict(return_components=True) returns (te, μ0, μ1). That holds for S/T/DR (yhat_cs=μ0, yhat_ts=μ1), but the X-learner returns (te, dhat_cs, dhat_ts) — two CATE estimates from the tau models, not potential outcomes (xlearner.py:214-215,244). Replaying the PR logic on synthetic_data(mode=1, n=100000):
- X-learner:
mean(mu0) == mean(mu1) == 0.689, somu1 - mu0 ≈ 0; residualsy - muare inflated → Γ=2 width 1.56 vs S/T ~1.26 (~24% off). - The Γ=1 point estimate still looks right (0.517 vs true 0.500) only because AIPW degenerates to IPW when
mu1≈mu0— which is exactly why the test passes.
Please extract μ0/μ1 learner-aware (for the X-learner use the outcome models models_mu_c/models_mu_t), or restrict + validate to S/T/DR and move the test off BaseXLearner. Adding an assertion that the Γ=1 bound ≈ the naive/true ATE would catch this class of error.
B2 — R-learner path raises an uncaught TypeError.
The description lists S/T/X/R/DR, but BaseRLearner.fit_predict() accepts no return_components, and both the try and except branches pass return_components=True, so the fallback doesn't rescue it:
TypeError: BaseRLearner.fit_predict() got an unexpected keyword argument 'return_components'
Please either drop R from the supported set (and the docstring) or raise a clear NotImplementedError naming the learner.
Non-blocking notes (sharpness wording — this is the conservative ZSBB/mean-residual bound rather than the quantile-balanced sharp Dorn-Guo bound; test seeding; propensity clipping) to follow separately.
Proposed changes
adds
SensitivityMSM, a newSensitivitysubclass implementing sensitivity bounds for the ATE under the Marginal Sensitivity Model (MSM). the existingSensitivitySelectionBiasonly works for linear outcome models, so there wasn't a way to get an actual numeric bound (rather than a refutation-style check) on how far the true ATE could move under unobserved confounding when using causalml's meta-learners (S/T/X/R/DR)MSM only needs propensity scores + the fitted outcome regression, both of which the existing meta-learners already produce, so no new dependency and no restriction to linear models
implementation is the closed-form, elementwise-sharp bound based on Dorn & Guo (2023, JASA) and the closed-form expressions in Dorn, Guo & Kallus (2024, arXiv:2112.11449). exposed as
get_msm_bounds(gamma=[...]), following the existingSensitivityclass conventions (sensitivity_estimate(), registration viaget_class_object/sensitivity_analysis())scope for this PR, per discussion in the issue...binary treatment only, ATE only.
plot_msm_bounds()andbreakdown_gamma()are deferred as follow-upsdocstring makes clear this is fragility diagnostic tooling (like the rest of the sensitivity module), not a green light for observational causal inference, and that the Gamma (propensity odds-ratio) parameterization is a different quantity from the partial-R² robustness value used by EconML/DoWhy
closes #916
Types of changes
What types of changes does your code introduce to CausalML?
Put an
xin the boxes that applyChecklist
Put an
xin the boxes that apply. You can also fill these out after creating the PR. If you're unsure about any of them, don't hesitate to ask. We're here to help! This is simply a reminder of what we are going to look for before merging your code.Further comments
If this is a relatively large or complex change, kick off the discussion by explaining why you chose the solution you did and what alternatives you considered, etc. This PR template is adopted from appium.