Technical Summary: Learning Who to Treat When Treatment is Missing

Problem Formulation

The paper addresses the challenge of learning optimal treatment allocation policies ($d: \mathcal{X} \to \{0, 1\}$) under budget constraints ($\kappa$) when treatment assignment data ($A$) is missing for a portion of the population. The objective is to maximize the expected utility $V(d) = E[Y(d(X))]$ by identifying individuals with the highest Conditional Average Treatment Effect (CATE), $\tau(X) = E[Y(1) - Y(0) | X]$.

Standard policy learning methods assume complete data. However, in observational studies (e.g., healthcare, social services), treatment records are frequently missing due to incomplete records, measurement errors, or data integration issues. The paper distinguishes between two missingness mechanisms:

1. Missing Completely Conditionally at Random (MCCAR): Missingness ($R=0$) depends only on observed covariates $X$. Formally, $R \perp\!\!\!\perp (A, Y) \mid X$.
2. Missing at Random (MAR): Missingness depends on observed covariates $X$ and observed outcomes $Y$. Formally, $R \perp\!\!\!\perp A \mid (X, Y)$.

Under MCCAR, standard complete-case analysis (discarding units with missing $A$) is valid but potentially inefficient. Under MAR, complete-case analysis is biased because the missingness mechanism is related to the outcome, violating the assumption that missingness is independent of potential outcomes given covariates. The paper restricts attention to these two settings, excluding Missing Not at Random (MNAR) scenarios where missingness depends on unobserved factors.

Methodology

The authors extend efficient influence function (IF) methods, originally developed for Average Treatment Effect (ATE) estimation under MAR, to the problems of Policy Value estimation and CATE estimation.

1. Identification Strategies

The paper proposes two identification strategies based on the missingness assumptions:

• Strategy One (MAR-based): Valid under both MCCAR and MAR. It utilizes all available data (including partially observed units) by reweighting observed outcomes. The policy value is identified as:
  $$V(d) = E\left[\frac{\beta_d(X)}{\gamma_d(X)}\right]$$
  where $\beta_d(X) = E[Y \lambda_d(X, Y) \mid X]$ and $\gamma_d(X) = E[\lambda_d(X, Y) \mid X]$. Here, $\lambda_d(X, Y) = P(A=d(X) \mid X, Y, R=1)$ is the propensity score conditioned on outcomes, and $\pi(X, Y) = P(R=1 \mid X, Y)$ is the probability of observing the treatment.
• Strategy Two (MCCAR-based): Valid only under MCCAR. It is a complete-case estimator using only units with observed treatments:
  $$V(d) = E[\nu_d(X)]$$
  where $\nu_d(X) = E[Y \mid X, A=d(X), R=1]$.

2. Efficient Estimators

The authors derive bias-corrected, doubly robust estimators using Influence Functions (IFs) to achieve semiparametric efficiency bounds.

• MCCAR Estimator ($\hat{\Phi}_{MCCAR}$): Based on the IF $\phi_{MCCAR}$, which relies on nuisance functions $\eta_d(X) = P(A=d(X), R=1 \mid X)$ and $\nu_d(X)$. It is consistent if either $\hat{\eta}_d$ or $\hat{\nu}_d$ is correctly specified.
• MAR Estimator ($\hat{\Psi}_{MAR}$): Based on the IF $\phi_{MAR}$, which leverages the full sample. It requires consistent estimation of nuisance functions $\pi(X, Y)$, $\lambda_d(X, Y)$, $\beta_d(X)$, and $\gamma_d(X)$. Crucially, it requires consistent estimation of $\gamma_d$ (the propensity score without outcome conditioning) in addition to either $\lambda_d$ or $\pi$.

3. Policy Learning (CATE Estimation)

To learn the optimal policy, the authors employ the DR-Learner framework. They construct pseudo-outcomes based on the difference between the MAR or MCCAR value function estimators for treatment ($A=1$) and control ($A=0$):

• $\tilde{\phi}_{MAR} = \hat{\phi}_{MAR}(1) - \hat{\phi}_{MAR}(0)$
• $\tilde{\phi}_{MCCAR} = \hat{\phi}_{MCCAR}(1) - \hat{\phi}_{MCCAR}(0)$

These pseudo-outcomes are regressed on covariates $X$ using generic machine learning methods (e.g., Random Forests) to estimate $\tau(X)$. The optimal policy is then defined as treating individuals where the estimated CATE exceeds a threshold determined by the budget constraint.

Key Contributions

1. Extension of Efficient Estimators: The authors extend efficient influence function methods from ATE estimation to policy value and CATE estimation under both MAR and MCCAR assumptions.
2. Theoretical Efficiency Proof: The paper proves that the MAR-based estimator is asymptotically more efficient than the MCCAR (complete-case) estimator whenever both are valid (i.e., under MCCAR assumptions). This provides a formal justification for preferring MAR-based methods even when the stronger MCCAR assumption holds, as the MAR estimator utilizes partially observed units to reduce variance.
3. Robustness to Misspecification: The proposed estimators possess double robustness properties. The MAR estimator remains consistent if either the outcome model ($\lambda_d$) or the missingness model ($\pi$) is correctly specified, provided the propensity score model ($\gamma_d$) is consistently estimated.
4. Empirical Validation: Comprehensive experiments on synthetic and semi-synthetic datasets (Voting and STAR) demonstrate that correctly specifying the missingness mechanism is critical. Misspecified estimators remain biased regardless of sample size, while the proposed MAR estimators achieve near-oracle performance when assumptions are satisfied.

Results

• Synthetic Experiments:
  • Convergence: In MAR settings, the MCCAR estimator exhibits persistent bias even with large sample sizes, whereas the MAR estimator converges to the oracle. In MCCAR settings, both estimators are unbiased, but the MAR estimator achieves slightly faster convergence rates.
  • Robustness: The MAR estimator is robust to misspecification of either the treatment propensity or the missingness mechanism, but not both simultaneously.
  • Efficiency Gains: Theoretical analysis and simulations confirm that the asymptotic relative efficiency (ARE) of the MAR estimator over the MCCAR estimator increases with the augmented propensity score, the missingness rate, and the correlation between them.
• Semi-Synthetic Experiments:
  • Using real-world data from randomized trials (Voting and STAR) with artificially induced missingness, the MAR-based estimators (DR-MAR) consistently outperformed complete-case methods and other baselines (IPW, Outcome Regression, MTRNET) in MAR settings, particularly when treatment data was severely limited (10-30% observed).
  • In MCCAR settings, most methods performed similarly, though the MAR-based approach remained competitive.
  • The results highlight that discarding incomplete observations (complete-case analysis) leads to suboptimal policies when missingness depends on outcomes.

Significance and Claims

The paper argues that complete-case analysis is not the efficient choice even when complete-case assumptions (MCCAR) hold. The authors claim that MAR-based estimators are theoretically superior because they:

1. Achieve lower asymptotic variance by leveraging partially observed units.
2. Require fewer identifying assumptions than MCCAR estimators (since MCCAR is a special case of MAR).
3. Provide a robust, nonparametric framework for policy learning that avoids the bias inherent in discarding data when missingness is outcome-dependent.

The authors conclude that in practice, where the missingness mechanism is often uncertain, defaulting to MAR-based estimation is the prudent strategy. Using MCCAR estimators under MAR assumptions leads to persistent bias, while using MAR estimators under MCCAR assumptions yields higher efficiency. This asymmetry reinforces the recommendation to adopt MAR-based methods for robust policy learning in the presence of missing treatment data.