Handling Missing Data in Downstream Tasks With Distribution-Preserving Guarantees

TL;DR

This paper introduces F3I, a novel distribution-preserving imputation method for missing data that improves downstream classification tasks, supported by theoretical guarantees and empirical validation across various applications.

Contribution

The paper proposes F3I, an innovative imputation approach with theoretical guarantees for distribution preservation, and demonstrates its effectiveness in high-dimensional, not-missing-at-random scenarios.

Findings

F3I outperforms existing imputation methods in accuracy.

F3I maintains data distribution better than traditional approaches.

F3I enhances downstream classification performance.

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.