Handling Missing Data in Downstream Tasks With Distribution-Preserving Guarantees

TL;DR

This work tackles missing data in downstream learning by introducing F3I, a fast distribution-preserving imputation method that iteratively refines a nearest-neighbor imputation through learned neighbor weights guided by a concave, differentiable objective. It provides theoretical guarantees on imputation quality and distribution preservation under MCAR, MAR, and MNAR, and extends the framework to joint training with downstream tasks via online learning and PCGrad-based gradient corrections. Theoretical results include high-probability MSE bounds and regret analyses for the online learner, with practical conditions for concavity and Lipschitz continuity of the objective. Empirically, F3I achieves competitive imputation performance and favorable runtimes across standard benchmarks and, when coupled with a classifier (PCGrad-F3I), demonstrates strong performance on joint imputation-classification tasks such as MNIST and drug-disease datasets, highlighting its practical impact for high-dimensional, real-world data.

Abstract

Missing feature values are a significant hurdle for downstream machine-learning tasks such as classification. However, imputation methods for classification might be time-consuming for high-dimensional data, and offer few theoretical guarantees on the preservation of the data distribution and imputation quality, especially for not-missing-at-random mechanisms. First, we propose an imputation approach named F3I based on the iterative improvement of a K-nearest neighbor imputation, where neighbor-specific weights are learned through the optimization of a novel concave, differentiable objective function related to the preservation of the data distribution on non-missing values. F3I can then be chained to and jointly trained with any classifier architecture. Second, we provide a theoretical analysis of imputation quality and data distribution preservation by F3I for several types of missing mechanisms. Finally, we demonstrate the superior performance of F3I on several imputation and classification tasks, with applications to drug repurposing and handwritten-digit recognition data.

Figures (24)

• Figure 1: Emprical validation of Theorem \ref{['thm:mse_bounds']} by comparing the value of the upper bound $N \times C^\text{miss}$ and $NF \times \mathcal{L}^\text{MSE}(X^t, X^\star)$ where $t$ is the final round for F3I. Left: MCAR setting. Center: MAR setting. Right: MNAR setting.
• Figure 2: $NF \times \mathcal{L}^\text{MSE}(X^t, X^\star)$ is linear in $\sigma^2$ regardless of the missingness mechanism (numerical values are reported in Table \ref{['tab:thm1']}).
• Figure 3: Evolution of the weight of each of the $K$-nearest neighbors for each sample as computed by F3I, where the $k$ neighbor is the $k^\text{th}$-nearest point, depending on the round $T$. Left: MCAR setting. Center: MAR setting. Right: MNAR setting.
• Figure 4: Cumulative regret for F3I and upper bound from Theorem \ref{['thm:regret_f3i']} across $100$ iterations. The blue points are always below the red lines. Left: MCAR setting. Center: MAR setting. Right: MNAR setting.
• Figure 5: Imputation on $2$ synthetic data sets $\times$$10$ different random seeds for generating missing values for F3I, K-nearest neighbor imputers troyanskaya2001missing (uniform or distance-based weights), mean imputation, MissForest stekhoven2012missforest, Optimal-Transport imputer muzellec2020missing and not-MIWAE ipsen2021not.

Theorems & Definitions (49)

• Definition 4.1
• Theorem 4.2
• Definition 4.3
• Theorem 4.4
• proof
• Theorem 5.1
• Remark B.7
• Proposition C.1
• proof
• Proposition C.2