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

📄 arXiv: 2412.19152v1 📥 PDF

作者: Jungkyu Kim, Kibok Lee, Taeyoung Park

分类: cs.LG, cs.AI, stat.ML

发布日期: 2024-12-26

🔗 代码/项目: GITHUB

💡 一句话要点

提出比例掩码自编码器PMAE，解决表格数据插补中掩码策略不当问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 表格数据插补 掩码自编码器 比例掩码 缺失数据 MLP 数据增强 机器学习

📋 核心要点

  1. 传统MAE在表格数据插补中采用的均匀随机掩码策略，忽略了表格数据固有的异构性和缺失模式，导致性能受限。
  2. PMAE通过计算并模拟数据集中真实的缺失比例，生成比例掩码，从而保留了缺失数据的原始分布，提升了插补的准确性。
  3. 实验表明，PMAE在多种缺失模式的表格数据集上表现出色，同时证明了MLP在表格数据处理中相比注意力机制的有效性和效率。

📝 摘要（中文）

本文针对表格数据插补任务，提出了比例掩码自编码器（PMAE）。现有的掩码自编码器（MAE）在表格数据上应用时，常用的均匀随机掩码策略会破坏缺失数据的分布，导致性能下降。为了解决这个问题，我们提出了一种比例掩码策略。该策略首先基于数据集中观察到的缺失比例计算缺失统计信息，然后生成与这些统计信息对齐的掩码，从而确保掩码后缺失数据的分布得以保留。此外，我们认为，简单的基于MLP的token混合方法在表格数据领域，相比于注意力机制，通常能提供更具竞争力的甚至更优越的性能，同时计算效率更高。实验结果验证了所提出的比例掩码策略在表格数据集中各种缺失数据模式下的有效性。

🔬 方法详解

问题定义：论文旨在解决表格数据插补任务中，由于数据缺失模式复杂且多样，传统MAE采用的均匀随机掩码策略无法有效利用数据中的缺失信息，导致插补性能下降的问题。现有方法没有充分考虑表格数据中缺失值的分布情况，直接应用随机掩码会破坏原始数据的缺失模式，从而影响模型的学习效果。

核心思路：论文的核心思路是设计一种与数据集中真实缺失模式相匹配的掩码策略，即比例掩码。通过分析数据集中缺失值的统计信息，生成与这些统计信息相符的掩码，从而在掩码过程中尽可能地保留原始数据的缺失分布。这样，模型在训练时能够更好地学习到数据中存在的缺失模式，从而提高插补的准确性。

技术框架：PMAE的整体框架仍然基于MAE，主要包含编码器和解码器两个部分。不同之处在于，PMAE在数据输入编码器之前，首先根据数据集的缺失统计信息生成比例掩码。编码器将掩码后的数据进行编码，解码器则根据编码后的信息重建原始数据。损失函数采用均方误差（MSE）来衡量重建数据与原始数据之间的差异。

关键创新：PMAE的关键创新在于提出了比例掩码策略。与传统的均匀随机掩码不同，比例掩码能够根据数据集中真实的缺失比例生成掩码，从而保留原始数据的缺失分布。此外，论文还验证了在表格数据处理中，简单的MLP结构在token混合方面可以取得与注意力机制相当甚至更好的效果，同时具有更高的计算效率。

关键设计：比例掩码的生成过程是关键。首先，计算数据集中每个特征的缺失比例。然后，根据这些缺失比例生成掩码，确保掩码后的数据具有与原始数据相似的缺失模式。论文中具体如何计算和应用这些比例信息，以及MLP的具体结构（层数、神经元数量等）等细节，需要在代码中进一步分析。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，PMAE在各种缺失模式的表格数据集上均取得了显著的性能提升。与传统的MAE相比，PMAE能够更准确地插补缺失值，尤其是在缺失比例较高的情况下。此外，论文还验证了MLP在表格数据处理中的有效性，证明了其在计算效率和性能方面都优于注意力机制。具体的性能提升幅度需要在论文中查找。

🎯 应用场景

PMAE可广泛应用于医疗、金融、生物信息等领域，这些领域的数据通常存在大量的缺失值。通过提高表格数据插补的准确性，PMAE可以帮助研究人员更好地分析和利用这些数据，从而为疾病诊断、风险评估、药物发现等应用提供更可靠的支持。未来，PMAE可以与其他数据增强技术相结合，进一步提升数据质量和模型性能。

📄 摘要（原文）

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