Treatment-Specific Outcome Predictions in AIPW

Why AIPW Needs Two Predictions per Unit

For each unit, the average treatment effect calculation uses two fitted outcome predictions: the expected outcome for that unit under treatment and the expected outcome for that same unit under control. These are conditional mean predictions, not two observed outcomes. Only one treatment condition is observed for each unit, but both fitted predictions contribute to the AIPW contrast.

The predictions provide the outcome-model part of the estimator. A residual correction then uses the observed outcome for the treatment condition the unit actually received. Keeping these pieces distinct is essential: the prediction under the unreceived condition is not observed, and the correction is not a substitute for predicting both conditions.

The target conditional means are \(m_a(x)=E[Y\mid A=a,X=x]\). Under consistency and conditional exchangeability, they identify the conditional means of the potential outcomes \(Y(1)\) and \(Y(0)\). The average treatment effect is the population average of their difference. This identification also requires positivity: in covariate strata represented in the target population, both treatment levels must have positive probability. These assumptions are about the causal interpretation; a fitted model by itself does not make them true.

Suppose the outcome model is a regression with treatment and covariate terms. To obtain the two predictions, retain a unit’s observed covariates, set the treatment input to 1, and calculate the fitted mean; then set it to 0 and calculate again. If separate models are fitted by treatment group, evaluate each group’s fitted model at that same unit’s covariates. In either approach, the two predictions must be on the outcome scale used in the AIPW formula.

Putting the Predictions into the AIPW Contrast

Let \(\hat e(X_i)\) be the supplied fitted propensity score, the estimated probability of treatment conditional on covariates. The unit-level AIPW contribution is the difference between the two predictions, plus a weighted residual for the treatment condition actually observed. Averaging these contributions gives the estimated average treatment effect.

For a treated unit, \(A_i=1\), so only the treated residual correction remains: \(\{Y_i-\hat m_1(X_i)\}/\hat e(X_i)\). For a control unit, \(A_i=0\), so only the control correction remains, with a minus sign: \(-\{Y_i-\hat m_0(X_i)\}/\{1-\hat e(X_i)\}\). The unobserved treatment-condition outcome does not appear as a residual; its fitted prediction still appears in the prediction difference.

Why the Residual Correction Uses These Predictions

The correction centers the observed outcome around the fitted mean for the condition the unit received. This is why the treated residual must use \(\hat m_1(X_i)\), while the control residual must use \(\hat m_0(X_i)\). A brief conditional-expectation calculation shows the alignment. Write \(e(x)=P(A=1\mid X=x)\) for the true propensity score and \(\mu_a(x)=E[Y\mid A=a,X=x]\) for the true conditional outcome mean.

If the outcome predictions are correct, so \(m_a(x)=\mu_a(x)\), the conditional mean of each observed residual is zero in its treatment group. For example, with a supplied score \(g(x)\) in the denominator, the treated correction has conditional expectation

The control correction is zero by the same reasoning when \(m_0(x)=\mu_0(x)\). The expected unit contribution is then \(\mu_1(x)-\mu_0(x)\). Conversely, if the supplied propensity score is correct, \(g(x)=e(x)\), the conditional means of the two weighted residuals are \(\mu_1(x)-m_1(x)\) and \(\mu_0(x)-m_0(x)\). Adding them with the AIPW signs to \(m_1(x)-m_0(x)\) gives \(\mu_1(x)-\mu_0(x)\). This explains the role of both predictions in the augmentation. The earlier material on double robustness develops the broader property.

Predictions Must Be on the Correct Outcome Scale

A fitted model may produce a linear predictor rather than a conditional mean on the outcome scale. For a binary outcome modeled with logistic regression, for instance, the linear predictor is converted to a predicted probability using the inverse-logit function. Substituting the linear predictor itself into the AIPW formula would mix scales and produce an invalid contrast.

A Complete Prediction Table and Its Corrections

A useful calculation workflow records both predictions for every unit before applying any residual correction. This makes it easier to verify that each prediction is evaluated at that unit’s covariates, that the observed outcome is paired with the correct treatment-specific prediction, and that the propensity-score denominator matches the observed treatment.

Diagnostics, Alternatives, and Common Errors

Before interpreting an AIPW result, inspect the predictions and their use in the formula. Diagnostics do not establish the causal assumptions, but they can reveal a calculation or model that is going wrong.

• Check the prediction scale. Predicted binary-outcome means should be probabilities, not logits. For a continuous outcome, the predictions and final contrast should use the outcome’s measurement units.
• Check treatment coding and model evaluation. Confirm that the treatment input is set to 1 and 0 in the right order, or that the correct separate model is used for each condition. Evaluate both predictions for every unit at that unit’s covariates.
• Check residual pairing and signs. Treated residuals use \(\hat m_1\) and denominator \(\hat e\). Control residuals use \(\hat m_0\), denominator \(1-\hat e\), and a minus sign in the AIPW expression.
• Inspect implausible predictions and extrapolation. Predictions far outside plausible outcome values can signal model misspecification or covariate patterns poorly represented in one treatment group. For binary outcomes, probabilities outside \([0,1]\) indicate a scale or implementation error.
• Inspect the supplied scores alongside predictions. Scores near 0 or 1 can make residual corrections large. Weak representation of one treatment condition at some covariate values also makes predictions there less secure. These are warning signs, not a guarantee that a particular estimate is biased or unbiased.

A common mistake is to use the observed outcome as the prediction for the treatment condition a unit received, then use the fitted model only for the unreceived condition. That reasoning confuses a realized outcome with a fitted conditional mean. The AIPW formula requires both fitted predictions; the observed outcome enters separately through the residual correction. Another error is to report the average prediction difference alone as the AIPW estimate. That is an outcome-regression plug-in estimate, not the augmented contrast shown here.

Outcome regression alone uses the fitted prediction difference and does not add these weighted residual corrections. Inverse-probability weighting relies on treatment weights rather than the two outcome predictions. AIPW combines a prediction contrast with residual corrections. Matching instead forms comparisons among units with similar covariates or scores. These approaches have different requirements and behavior; none removes the need to assess whether the causal assumptions and data support the target comparison. The AIPW calculation also cannot make unsupported predictions reliable when treatment overlap is poor.