Not Another Imputation Method: A Transformer-based Model for Missing Values in Tabular Datasets

TL;DR

NAIM is a transformer-based model that directly handles missing values in tabular data without imputation, improving predictive accuracy and robustness over existing methods.

Contribution

The paper introduces NAIM, a novel transformer model with feature-specific embeddings and masked self-attention to effectively learn from incomplete data without imputation.

Findings

NAIM outperforms state-of-the-art models on multiple datasets.

NAIM demonstrates superior resilience to missing data.

The model's regularization enhances generalization from incomplete data.

Abstract

Handling missing values in tabular datasets presents a significant challenge in training and testing artificial intelligence models, an issue usually addressed using imputation techniques. Here we introduce "Not Another Imputation Method" (NAIM), a novel transformer-based model specifically designed to address this issue without the need for traditional imputation techniques. NAIM's ability to avoid the necessity of imputing missing values and to effectively learn from available data relies on two main techniques: the use of feature-specific embeddings to encode both categorical and numerical features also handling missing inputs; the modification of the masked self-attention mechanism to completely mask out the contributions of missing data. Additionally, a novel regularization technique is introduced to enhance the model's generalization capability from incomplete data. We extensively evaluated NAIM on 5 publicly available tabular datasets, demonstrating its superior performance over 6 state-of-the-art machine learning models and 5 deep learning models, each paired with 3 different imputation techniques when necessary. The results highlight the efficacy of NAIM in improving predictive performance and resilience in the presence of missing data. To facilitate further research and practical application in handling missing data without traditional imputation methods, we made the code for NAIM available at https://github.com/cosbidev/NAIM.