# Distribution Shift in Missing Data Imputation: A Risk-Based Perspective and Importance-Weighted Correction under MAR

arXiv: https://arxiv.org/abs/2602.06713 · HTML: https://mlsanity.com/paper/2602.06713/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 6 of its 22 references that are in the map.
- First posted on arXiv 6 Feb 2026.

## Builds on (most cited first, 6 of 6)

- [GAIN: Missing Data Imputation using Generative Adversarial Nets](https://mlsanity.com/paper/1806.02920/) (2018)
- [Missing Data Imputation using Optimal Transport](https://mlsanity.com/paper/2002.03860/) (2020)
- [HyperImpute: Generalized Iterative Imputation with Automatic Model Selection](https://mlsanity.com/paper/2206.07769/) (2022)
- [Doubly Robust Calibration of Prediction Sets under Covariate Shift](https://mlsanity.com/paper/2203.01761/) (2022)
- [Domain Adaptation under Missingness Shift](https://mlsanity.com/paper/2211.02093/) (2022)
- [Sharing pattern submodels for prediction with missing values](https://mlsanity.com/paper/2206.11161/) (2022)

## Built on by (most cited first, 1 of 1)

- [Shapley Value Estimation for Multi-Site Data with Blockwise-Missing Features](https://mlsanity.com/paper/2609.14902/) (2026)

## Abstract (arXiv)

Missing data imputation, where a model is trained on observed data to estimate unobserved values, is a fundamental problem in machine learning. In this paper, we rigorously formulate imputation model learning as a mean-squared error risk minimisation problem. We show that when the probability of missingness depends on the data, many state-of-the-art methods fail to account for the resulting distribution shift between the observed data used for training and the full data distribution used for evaluation. Consequently, these approaches do not minimise mean-squared error on the full data distribution. Instead, we propose a novel imputation algorithm designed to learn an imputation model from the observed data while explicitly accounting for this distribution shift. Simulation studies show consistent improvements over otherwise identical uncorrected baselines, with average reductions of 3% in RMSE and 7% in Wasserstein distance.