Enhanced multi-task learning of imputation and prediction via feature relationship graph learning
Keywords: Missing values, imputation, feature selection, graph nueral network, multi-task learning
Abstract: Missing values present significant challenges in machine learning, often degrading predictive performance. Traditional and deep learning imputation methods often overlook the relationships between features and their connections to downstream tasks. To address these gaps, we propose PIG (multi-task learning of Prediction and Imputation via feature-relationship Graph learning), a model that integrates imputation and prediction by leveraging feature interdependencies.
PIG utilizes a graph-based approach to capture intricate feature relationships, thereby enhancing the accuracy of both imputation and downstream tasks. Our strategic training process begins with pre-training for both tasks, ensuring the model learns effective representations. This is followed by fine-tuning the entire model to further optimize imputation and downstream tasks simultaneously.
We evaluated our method using nine benchmark datasets, three for regression and six for classification.
Our method showed superior imputation and prediction performance across nine datasets, achieving an average rank of 1.33 for both imputation and regression tasks and 1.83 for imputation and 1.17 for classification tasks. Additionally, in sensitivity analysis with respect to missing rates, our method demonstrated its robustness, especially in predictive performance, compared to other methods that showed significant degradation.
Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 6767
Loading