Technical Summary: Distribution Shift in Missing Data Imputation

Problem Formulation

The paper addresses the fundamental challenge of missing data imputation, specifically focusing on the discrepancy between the distribution of observed data used for training and the full data distribution required for evaluation. While many state-of-the-art imputation methods (e.g., GAIN, MIWAE, MICE, MissForest) operate under a risk minimisation framework, they typically train on observed entries without explicitly correcting for the distributional shift induced by the missingness mechanism.

The authors formalise the problem as minimising the Mean Squared Error (MSE) risk over the full data distribution. They demonstrate that when the probability of missingness depends on the data (specifically under the Missing At Random (MAR) assumption), the training distribution (observed data) and the target distribution (full data) differ. This creates a covariate shift problem where standard empirical risk minimisation on observed data yields biased estimators for the full data risk.

Methodology

Theoretical Foundation

The core theoretical contribution is the derivation of the imputation risk as an importance-weighted expectation over the observed data.

• Risk Decomposition: The global MSE risk is decomposed into a sum of coordinate-wise risks. For a specific coordinate $i$ with missingness, the risk $J_i(g_i)$ is defined over the unobserved subpopulation ($R_i=0$).
• Covariate Shift Connection: Under the MAR assumption ($p(r|x) = p(r|x_{obs})$) and an overlap condition, the authors show that the risk over the unobserved subpopulation can be rewritten as an expectation over the observed subpopulation ($R_i=1$) weighted by an importance factor:
  $$w_i(x_{obs}, r_{\neg i}) = \frac{p(x_{obs}, r_{\neg i} | R_i = 0)}{p(x_{obs}, r_{\neg i} | R_i = 1)}$$
  This weight corrects for the shift in the marginal distribution of the covariates between the observed and unobserved groups. Existing methods are shown to implicitly omit this weighting term, effectively optimising for the wrong distribution.

Proposed Algorithm: Weighted Iterative Imputation

To operationalise this theory, the authors propose a Weighted Iterative Imputation algorithm (Algorithm 1), which extends the standard round-robin iterative approach (e.g., MICE, HyperImpute).

1. Iterative Framework: The algorithm cycles through partially observed variables. In each step, a conditional imputation model $g_i$ is fitted for the target variable using the current values of all other variables as predictors.
2. Surrogate Risk: Instead of minimising the standard reconstruction loss, the algorithm minimises a surrogate risk that incorporates estimated importance weights:
   $$J_i^{(t)}(g_i) = \mathbb{E}\left[ w_i(\hat{X}_{\neg i}^{(t)}) (g_i(\hat{X}_{\neg i}^{(t)}) - X_i)^2 \mid R_i = 1 \right]$$
3. Weight Estimation: Since the true weights depend on unknown distributions, they are estimated at each iteration (Algorithm 2) using a binary classification model (logistic regression) to estimate the missingness probability $p(R_i=1 | \hat{X}_{\neg i})$. The weights are then computed as the inverse probability ratio.
4. Practical Stabilisation: To manage the high variance often associated with importance weighting, the authors employ several stabilisation techniques:
   • Clipping weights at the 5th and 95th percentiles.
   • Removing features with near-zero variance.
   • Introducing a tempering parameter $\gamma$ (where weights become $w_i^\gamma$) to control the influence of the weights, tuned via cross-validation.
   • A "warm-starting" strategy for tuning $\gamma$ to reduce computational overhead.

Key Contributions

1. Theoretical Formulation: The paper rigorously establishes that imputation under MAR is equivalent to a collection of covariate shift problems, deriving the specific importance weights required to align the observed data distribution with the full data distribution.
2. Novel Algorithm: It introduces a practical, round-robin iterative algorithm that integrates these importance weights into the training objective, explicitly correcting for the distributional shift.
3. Empirical Validation: The method is benchmarked against unweighted baselines (using identical conditional model classes) and state-of-the-art methods across diverse datasets (tabular, time-series, and image data) and model classes (Linear, Random Forest, MLP).

Experimental Results

The authors evaluated their method on eight datasets, simulating MAR conditions with varying strengths of dependence between missingness and observed covariates.

• Performance Gains: The weighted approach consistently reduced the Root Mean Squared Error (RMSE) and Wasserstein distance (WD) compared to unweighted baselines.
  • Average reduction of 3% in RMSE and 7% in Wasserstein distance.
  • Improvements were most pronounced for simpler models (Linear) and in settings with strong covariate shift (high MAR dependence).
• Model Expressivity: Theoretical analysis and empirical results indicate that the benefit of weighting diminishes as the expressivity of the conditional model increases (e.g., Random Forests and MLPs). This is attributed to the bias-variance trade-off; highly expressive models can approximate the target function well even without explicit weighting, whereas simpler models rely more heavily on the correction.
• Downstream Tasks: In downstream regression and classification tasks, the weighted imputation generally led to statistically significant improvements in predictive performance, particularly for linear models and when compared to unimputed baselines.
• Robustness: The method showed gradual performance degradation under violations of the MAR assumption (MNAR) or weight model misspecification, rather than abrupt failure, provided the primary drivers of missingness were captured in the observed data.
• Efficiency: The computational cost was manageable, with runtimes increasing by only ~10% for complex models (RF/MLP) and showing significant speedups for linear models due to the specific implementation details.

Significance and Claims

The paper claims that distributional shift is an inherent and often overlooked issue in risk-based imputation. By framing imputation as a covariate shift problem, the authors provide a principled justification for using importance weighting.

• The significance lies in demonstrating that explicitly accounting for the shift between observed and full data distributions leads to more accurate imputation and better preservation of the underlying data distribution (fidelity).
• The authors position their work as a necessary refinement to existing risk-minimisation frameworks, arguing that without this correction, imputation models are optimised for the observed data distribution rather than the true data distribution required for valid inference.
• The paper remains modest regarding limitations, acknowledging that the method may offer limited gains in weak-MAR regimes or with highly expressive models where the approximation error is already low, and that performance depends on the quality of the weight estimation.