# A Bayesian marginal structural model (IPW) in a single model

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(
"y_obs",
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

• 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? Commented Mar 13, 2023 at 0:44
• 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. Commented Mar 13, 2023 at 7:06