Skip to content

Add Sensitivity bounds for ATE via Marginal Sensitivity Model#925

Open
aman-coder03 wants to merge 1 commit into
uber:masterfrom
aman-coder03:feature/sensitivity-msm
Open

Add Sensitivity bounds for ATE via Marginal Sensitivity Model#925
aman-coder03 wants to merge 1 commit into
uber:masterfrom
aman-coder03:feature/sensitivity-msm

Conversation

@aman-coder03

@aman-coder03 aman-coder03 commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Proposed changes

adds SensitivityMSM, a new Sensitivity subclass implementing sensitivity bounds for the ATE under the Marginal Sensitivity Model (MSM). the existing SensitivitySelectionBias only 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 existing Sensitivity class conventions (sensitivity_estimate(), registration via get_class_object/sensitivity_analysis())

scope for this PR, per discussion in the issue...binary treatment only, ATE only. plot_msm_bounds() and breakdown_gamma() are deferred as follow-ups

docstring 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 x in the boxes that apply

  • Bugfix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Documentation Update (if none of the other choices apply)

Checklist

Put an x in 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.

  • I have read the CONTRIBUTING doc
  • I have signed the CLA
  • Lint and unit tests pass locally with my changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have added necessary documentation (if appropriate)
  • Any dependent changes have been merged and published in downstream modules

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.

@jeongyoonlee jeongyoonlee left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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, so mu1 - mu0 ≈ 0; residuals y - mu are 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.

@jeongyoonlee jeongyoonlee added the enhancement New feature or request label Jul 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Sensitivity bounds for CATE via Marginal Sensitivity Model (MSM)

2 participants