AIPW

Cross-fit augmented inverse-probability weighting

from _api_doc_utils import *

1 Where it fits

Group: Causal inference

AIPW estimates a binary-treatment ATE by combining outcome regressions and a propensity model:

\[ \hat\tau = n^{-1}\sum_i \left[\hat\mu_1(x_i)-\hat\mu_0(x_i) + \frac{d_i(y_i-\hat\mu_1(x_i))}{\hat e(x_i)} - \frac{(1-d_i)(y_i-\hat\mu_0(x_i))}{1-\hat e(x_i)}\right]. \]

The nuisance functions are cross-fit ridge models. Fold assignment is deterministically shuffled within treatment strata, preventing avoidable single-arm training folds while preserving seed reproducibility.

2 Cross-fitted nuisance models and ATE

The class uses seeded treatment-stratified outer folds. In each training complement it fits three unpenalized-intercept ridge regressions:

\[ \hat\mu_1(x)\ \text{from treated outcomes},\qquad \hat\mu_0(x)\ \text{from control outcomes},\qquad \hat\pi(x)\ \text{from a linear regression of }D\text{ on }X. \]

Thus the propensity model is a ridge linear-probability model, not logistic regression. Its held-out predictions are clipped to \([c,1-c]\). Each nuisance objective is

\[ \sum_{i\in I_{\mathrm{train}}}(r_i-a-X_i'b)^2+\lambda\|b\|_2^2, \]

with raw, unstandardized features. A penalty grid is selected independently for each nuisance and outer fold using deterministic inner-fold MSE.

The held-out augmented inverse-probability pseudo-outcome is

\[ \phi_i = \hat\mu_1(X_i)-\hat\mu_0(X_i) +\frac{D_i\{Y_i-\hat\mu_1(X_i)\}}{\hat\pi(X_i)} -\frac{(1-D_i)\{Y_i-\hat\mu_0(X_i)\}}{1-\hat\pi(X_i)}, \]

and the estimate is \(\hat\tau=E_n[\phi_i]\).

3 Inference

The influence score used for covariance is \(\phi_i-\hat\tau\). Vanilla is its uncorrected iid empirical variance divided by \(n\); HC1, Bartlett Newey-West, and cluster options apply the common finite-sample corrections. The calculation is conditional on the realized cross-fit and selected nuisance penalties. There is no bootstrap, repeated sample splitting, or built-in Wald method.

The usual double-robust interpretation requires at least one nuisance structure to be correctly specified and standard overlap and sampling assumptions. Clipping stabilizes denominators but changes the estimating equation and does not diagnose lack of overlap. Every outer training fold must contain treated and control observations.

4 Performance and numerical behavior

Each outer fold fits separate treated and control outcome models plus a propensity model. With \(K\) folds, \(L\) candidate penalties, and \(F\) inner folds, the workload is roughly \(3KLF\) dense ridge solves, with smaller outcome samples inside treatment arms. No factorization is reused. The class stores all held-out nuisance predictions and selected penalties. Small treatment groups, many controls, raw feature scaling, and predictions frequently hitting the clipping bounds are the main stability concerns.

5 Python API

Constructor: cm.AIPW

Call fit(y, d, x) with binary treatment d. summary() reports ate, se, vcov, and selected penalties for the outcome and propensity nuisance models.

print(inspect.signature(cm.AIPW))

(penalty=None, cv=5, n_folds=5, propensity_clip=0.02, seed=42)

cls = cm.AIPW
display(HTML(html_table(["Public method"], public_methods(cls))))

Public method
fit(self, /, y, d, x)
summary(self, /, vcov=None, lags=None, clusters=None)

6 Minimal example

rng = np.random.default_rng(14)
x = rng.normal(size=(420, 3))
pi = 1 / (1 + np.exp(-(0.1 + x @ np.array([0.6, -0.3, 0.2]))))
d = rng.binomial(1, pi, size=420).astype(float)
y = 0.5 + x @ np.array([0.2, -0.1, 0.3]) + 1.0 * d + rng.normal(size=420)
model = cm.AIPW(penalty=np.logspace(-4, 1, 10), cv=3, n_folds=4, seed=2)
model.fit(y, d, x)
print(model.summary()['ate'])
print(model.summary()['se'])

1.0852907049864402
0.10693192038330385

7 summary() contract

The table below is generated by fitting the live class in this repository and then inspecting summary(). Shapes are shown because most values are plain NumPy arrays or scalars.

rng = np.random.default_rng(114)
x = rng.normal(size=(160, 3))
pi = 1 / (1 + np.exp(-(0.1 + x @ np.array([0.6, -0.3, 0.2]))))
d = rng.binomial(1, pi, size=160).astype(float)
y = 0.5 + x @ np.array([0.2, -0.1, 0.3]) + d + rng.normal(size=160)
model = cm.AIPW(penalty=np.logspace(-4, 1, 6), cv=3, n_folds=4, seed=2)
model.fit(y, d, x)
summary = model.summary()
display(HTML(html_table(["summary() key", "shape"], summary_shape_rows(summary))))

summary() key	shape
ate	()
se	()
vcov	(1, 1)
outcome0_penalties	(4,)
outcome1_penalties	(4,)
propensity_penalties	(4,)