Handling Missing Data in Downstream Tasks With Distribution-Preserving Guarantees
This addresses missing data challenges in machine learning for applications such as healthcare and image recognition, offering a novel imputation approach with theoretical guarantees.
The paper tackles the problem of missing feature values in downstream classification tasks by proposing F3I, an imputation method that preserves data distribution and improves classification performance, demonstrating superior results on tasks like drug repurposing and handwritten-digit recognition.
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.