To Predict or Not To Predict? Proportionally Masked Autoencoders for Tabular Data Imputation

TL;DR

The paper tackles imputation for heterogeneous tabular data, where uniform masking in masked autoencoders can misalign with column-wise missingness. It introduces Proportionally Masked AutoEncoder PMAE, which uses a logit-based masking function M_j^{PM} driven by observed proportions p_obs to balance prediction and reconstruction across columns, and advocates MLP-Mixer token mixing over Transformers for tabular data. Empirical results across nine real-world datasets and multiple missing patterns show PMAE, especially PMAE-mix, achieving state-of-the-art imputation accuracy and strong downstream-task performance, with gains up to 34.1% under challenging MNAR General patterns. A unified Imputation Accuracy metric combines categorical accuracy and numerical R^2, enabling a holistic assessment of imputation quality. The work demonstrates that proportional masking and simpler yet effective MLP-based mixing yield robust improvements, offering practical implications for real-world tabular data imputation under diverse missingness scenarios.

Abstract

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}.

Figures (8)

• Figure 1: The simplest missing pattern is Monotone, where some columns are fully observed, common in longitudinal studies. As the missing pattern becomes General (with all columns prone to missing with varying ratios) imputation becomes more challenging. Practitioners need strategies to address these patterns, which are currently under explored.
• Figure 2: PMAE. Given the observed mask $\mathbf{m}$, we calculate the observed proportions and apply an additional mask, $\mathbf{m}^-$, where the extra masking probabilities are inversely proportional to the observed proportions.
• Figure 3: Masking data inversely to observed proportions, prioritizing prediction when data is sparse and reconstruction when more data is observed (the last column of (b)).
• Figure 4: Imputation accuracy of state-of-the-art methods across 9 benchmark datasets on missing patterns (Monotone, Quasi Monotone, and General) under NMAR mechanism. Methods are arranged such that the most recent is on the left.
• Figure 5: Magnitude comparison of representations after token mixing via. Self-Attention vs. MLP on Shoppers dataset. Absolute value is applied for the sake of analysis.