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

TL;DR

This paper introduces a proportional masking strategy for masked autoencoders in tabular data imputation, improving performance by preserving missingness distribution and demonstrating that simple MLP-based mixing can outperform attention mechanisms.

Contribution

It proposes a novel proportional masking approach for MAEs in tabular data and shows that simple MLP-based token mixing can be more effective and efficient than attention mechanisms.

Findings

Proportional masking improves imputation accuracy across various missing data patterns.

Simple MLP-based token mixing often outperforms attention mechanisms in tabular data.

Experimental results validate the effectiveness of the proposed masking strategy.

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