Inspired by Richard McElreath's "Full Luxury Bayes" in his Statistical Rethinking course, I wanted to implement a "Full Luxury Bayesian Marginal Structural Model".

Briefly: MSMs are a two-step model for the average treatment effect. First, you regress the (binary) treatment $A$ on confounders $X$; second, you compute the inverse probability weights $w_i=\frac{1}{\Pr[A=a_i|X=X_i]}$; and third, regress the outcome on the treatment ($Y \sim 1+A$) weighted by $w$.

I thought this is a classic case for this type of multiple submodels within a single "full" model, since there is one regression, a deterministic computation and a second regression. Very similar to the example McElreath presents.
Unfortunately, when I experiment, I can't recover the true parameters, and I believe it's not a software bug (see details below), but an actual consequence of learning the propensity and outcome models jointly (thus posting here and not on stackoverflow).

I wonder if someone could explain why this goes wrong.

The model that doesn't work:

import pymc as pm

with pm.Model() as msm_model:
    # Treatment model:
    intercept_a = pm.Normal("intercept_a", mu=0, sigma=2)
    betas_a = pm.Normal("betas_a", mu=0, sigma=2, shape=X.shape[1])
    mu_lin_a = pm.Deterministic(
        "mu_lin_a", intercept_a + pm.math.dot(X, betas_a),
    p_a1 = pm.Deterministic("p_a1", pm.math.sigmoid(mu_lin_a))  # Pr[A=1|X]
    a_obs = pm.Bernoulli("a_obs", p_a1, observed=a)
    p_a0 = pm.Deterministic("p_a0", 1-p_a1)  # Pr[A=0|X]
    p_a = pm.Deterministic("p_a", (a*p_a1) + (1-a)*p_a0)  # Pr[A=a_i|X]
    ipa = pm.Deterministic("ipa", 1/p_a)
    # Outcome model (MSM):
    intercept_y = pm.Normal("intercept_y", mu=0, sigma=3)
    betas_y = pm.Normal("betas_y", mu=0, sigma=3)
    sigma_y = pm.HalfNormal("sigma_y", sigma=3)
    mu_lin_y = pm.Deterministic(
        "mu_lin_y", intercept_y + betas_y*a,
    # This is how to define a weighted regression in PyMC:
    y_obs = pm.Potential(
        ipa * pm.logp(pm.Normal.dist(mu=mu_lin_y, sigma=sigma_y), y)

The reason I think there isn't a bug in this model is that I can make it work by changing two things separately:

  1. If I precompute the IP-weights beforehand using regular logistic regression (say, in statsmodels).
  2. If instead of a weighted outcome regression I compute the weighted average in each group (Horvitz–Thompson estimator).
    Furthermore, when I do so the (averaged over chains) propensities (p_a1) suddenly match the ones I get from the non-Bayesian model, whereas they do not match in the "full" model.

This why I think there's some inherent fault in the joint model which I don't understand.

Sample data:

import numpy as np

def generate_data(seed=0, N=1000, D=1, effect=0):
    rng = np.random.default_rng(seed)
    X = rng.normal(0, 1, size=(N, D))
    beta_xa = rng.normal(2, 1, size=D)  # 3.184
    a_logit = 1 + X@beta_xa + rng.normal(0, 0, size=N)
    a_propensity = 1 / (1 + np.exp(-a_logit))
    a = rng.binomial(1, a_propensity)
    beta_xy = rng.normal(-2, 1, size=D)  # -1.418
    y = 1 + X@beta_xy + a*effect + rng.normal(0, 1, size=N)

    return X, a, y
  • $\begingroup$ IPW approaches can give pretty erratic estimates, in particular when weights are close to 0 or 1. Might this be the case for certain $(A,X)$ pairs? If you clip your weights to some interval like [.01, 1 - .01] do the results get any better? $\endgroup$
    – Daniel
    Commented Mar 13, 2023 at 0:44
  • $\begingroup$ Thanks, @Daniel, I tried it now just to be on the safe side, but that's not it. IPW can indeed have large variance, but I get a pretty large bias. This is well-behaved data with a well-specified model. The fact that I can recover the ATE if I use weighted averages (within the model, using deterministic pm.math.sum() operations) instead of a second regression model, suggests it is something with the joint fitting of the model. $\endgroup$
    – ehudk
    Commented Mar 13, 2023 at 7:06


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.