To Predict or Not To Predict? Proportionally Masked Autoencoders for Tabular Data Imputation
This work addresses imputation challenges for tabular data users, but it is incremental as it builds on existing MAE methods with a modified masking strategy.
The paper tackles the problem of suboptimal performance in masked autoencoders for tabular data imputation due to uniform random masking disrupting missingness distribution, and proposes a proportional masking strategy that preserves this distribution, validated with experimental results across various missing data patterns.
Masked autoencoders (MAEs) have recently demonstrated effectiveness in tabular data imputation. However, due to the inherent heterogeneity of tabular data, the uniform random masking strategy commonly used in MAEs can disrupt the distribution of missingness, leading to suboptimal performance. To address this, we propose a proportional masking strategy for MAEs. Specifically, we first compute the statistics of missingness based on the observed proportions in the dataset, and then generate masks that align with these statistics, ensuring that the distribution of missingness is preserved after masking. Furthermore, we argue that simple MLP-based token mixing offers competitive or often superior performance compared to attention mechanisms while being more computationally efficient, especially in the tabular domain with the inherent heterogeneity. Experimental results validate the effectiveness of the proposed proportional masking strategy across various missing data patterns in tabular datasets. Code is available at: \url{https://github.com/normal-kim/PMAE}.