2025-11-15T15:43:11.510797

TabImpute: Accurate and Fast Zero-Shot Missing-Data Imputation with a Pre-Trained Transformer

Feitelberg, Saha, Choi et al.
Missing data is a pervasive problem in tabular settings. Existing solutions range from simple averaging to complex generative adversarial networks. However, due to huge variance in performance across real-world domains and time-consuming hyperparameter tuning, no default imputation method exists. Building on TabPFN, a recent tabular foundation model for supervised learning, we propose TabImpute, a pre-trained transformer that delivers accurate and fast zero-shot imputations requiring no fitting or hyperparameter tuning at inference-time. To train and evaluate TabImpute, we introduce (i) an entry-wise featurization for tabular settings, which enables a $100\times$ speedup over the previous TabPFN imputation method, (ii) a synthetic training data generation pipeline incorporating realistic missingness patterns, which boosts test-time performance, and (iii) MissBench, a comprehensive benchmark for evaluation of imputation methods with $42$ OpenML datasets and $13$ missingness patterns. MissBench spans domains such as medicine, finance, and engineering, showcasing TabImpute's robust performance compared to $11$ established imputation methods.
academic

TabImpute: Accurate and Fast Zero-Shot Missing-Data Imputation with a Pre-Trained Transformer

基本信息

  • 论文ID: 2510.02625
  • 标题: TabImpute: Accurate and Fast Zero-Shot Missing-Data Imputation with a Pre-Trained Transformer
  • 作者: Jacob Feitelberg, Dwaipayan Saha, Kyuseong Choi, Zaid Ahmad, Anish Agarwal, Raaz Dwivedi
  • 分类: cs.LG (Machine Learning)
  • 发表时间: 2025年10月 (Preprint. Under Review)
  • 论文链接: https://arxiv.org/abs/2510.02625v2

摘要

缺失数据是表格数据中的普遍问题。现有解决方案从简单的均值填充到复杂的生成对抗网络不等。然而，由于在真实世界领域中性能差异巨大且超参数调优耗时，目前并不存在默认的插补方法。基于TabPFN（一个最新的表格监督学习基础模型），本文提出TabImpute，这是一个预训练的Transformer，能够在推理时提供准确且快速的零样本插补，无需拟合或超参数调优。为了训练和评估TabImpute，作者引入了：(i) 针对表格设置的条目级特征化，相比之前的TabPFN插补方法实现了100倍加速；(ii) 结合现实缺失模式的合成训练数据生成管道，提升测试时性能；(iii) MissBench，一个包含42个OpenML数据集和13种缺失模式的综合插补方法评估基准。MissBench涵盖医学、金融和工程等领域，展示了TabImpute相比11种成熟插补方法的稳健性能。

研究背景与动机

问题定义

缺失数据在表格数据中无处不在，影响统计学家、经济学家、卫生官员和企业。例如，医疗数据集可能缺少血压测量记录，或从多个来源合并的数据集可能只共享部分特征。无论来源如何，在使用统计或机器学习模型之前，必须将缺失数据插补为数值。

问题重要性

  1. 普遍性：缺失数据在各个领域都是常见问题
  2. 必要性：大多数机器学习算法无法直接处理缺失值
  3. 复杂性：不同的缺失机制需要不同的处理策略

现有方法局限性

  1. 性能差异大：现有方法在不同领域和数据集上表现差异巨大
  2. 超参数调优：需要耗时的超参数调整过程
  3. 缺乏通用方法：没有默认的插补方法适用于所有场景
  4. 特定场景限制：每种方法通常针对特定设置设计

研究动机

基于TabPFN在表格监督学习中的成功，作者希望开发一个能够：

  1. 实现零样本插补（无需训练或调参）
  2. 在多种缺失模式下表现稳健
  3. 提供快速准确的插补结果
  4. 适用于不同领域的表格数据

核心贡献

  1. 提出TabImpute模型：基于TabPFN架构的预训练Transformer，实现准确快速的零样本缺失数据插补
  2. 创新的条目级特征化（Entry-wise Featurization）：相比TabPFN的列级插补方法实现了100倍加速
  3. 综合的合成数据生成管道：包含13种现实缺失模式的训练数据生成方法，提升模型泛化能力
  4. 构建MissBench基准：包含42个OpenML数据集和13种缺失模式的综合评估基准
  5. TabImpute+集成方法：通过自适应权重集成TabImpute和EWF-TabPFN，达到最佳性能

方法详解

任务定义

给定带有缺失值的表格数据矩阵X，其中X*为完整矩阵，Ω为缺失条目索引集合，目标是预测所有缺失条目的值。

模型架构

1. 条目级特征化（Entry-wise Featurization, EWF）

传统方法采用列级插补，而本文提出条目级方法：

  • 对每个条目(i,j)，构造特征向量：(i ⊕ j ⊕ Xi,: ⊕ X:,j)
  • 其中Xi,:表示第i行，X:,j表示第j列，⊕表示拼接
  • 目标值为yij = X*ij
  • 创建大小为nm × (n+m)的特征矩阵

2. 架构修改

基于TabPFN架构，做出一个关键修改：

  • 移除注意力掩码，允许训练点关注测试点
  • 这样做的原因是测试集使用已观察到的数据创建，不存在数据泄露问题

3. 合成训练数据生成

数据生成：使用线性因子模型（Linear Factor Models）

Y = UV^T

其中U ∈ R^(m×k)，V ∈ R^(n×k)，k ≪ n,m

缺失模式：实现13种缺失模式

  • 1种MCAR（Missing Completely At Random）
  • 1种MAR（Missing At Random）
  • 11种MNAR（Missing Not At Random）

4. 多模式训练

使用自适应算法确定每个批次中缺失模式的比例：

  • 每s个梯度步骤重新计算比例
  • 对各模式的损失值应用softmax
  • 自适应降低表现好的模式权重，提高表现差的模式权重

技术创新点

  1. 并行化处理：条目级特征化使得可以并行预测所有缺失值，而非逐列处理
  2. 零样本能力：预训练模型无需在目标数据上微调即可直接使用
  3. 多模式适应：通过自适应训练策略处理多种缺失模式
  4. 集成策略：TabImpute+通过最优权重集成不同方法的优势

实验设置

数据集

MissBench基准包含：

  • 42个OpenML数据集
  • 涵盖医学、工程、教育等多个领域
  • 数据集大小从50×5到170×55不等
  • 仅包含数值特征且原始无缺失值的数据集

评价指标

插补准确度：

  1. 计算每种方法的RMSE：1∣Ω∣∑(i,j)∈Ω(Xijtrue−Xijimputed)2\sqrt{\frac{1}{|\Omega|}\sum_{(i,j)\in\Omega}(X^{true}_{ij} - X^{imputed}_{ij})^2}
  2. 在每个任务内进行最小-最大归一化
  3. 插补准确度 = 1 - 归一化RMSE

对比方法

11种成熟的插补方法：

  • 列均值插补
  • SoftImpute
  • MissForest
  • ICE/MICE
  • GAIN
  • MIWAE
  • 最优传输方法
  • K近邻
  • HyperImpute
  • TabPFN原始插补方法

实现细节

  • 训练：8个H200 GPU，约一周时间
  • 处理2500万个合成表格
  • 学习率：0.0001，批次大小：64
  • 自适应更新间隔：s=50步

实验结果

主要结果

整体性能（表1）：

  • TabImpute+：0.833 ± 0.213（最佳）
  • HyperImpute：0.766 ± 0.259
  • 最优传输：0.765 ± 0.227
  • MissForest：0.754 ± 0.248

运行时间（图1b）：

  • TabImpute在GPU上运行时间最短
  • 相比TabPFN实现了显著加速
  • CPU版本仍保持竞争力

不同缺失模式下的表现

TabImpute+在几乎所有缺失模式下都取得最佳性能：

  • NN-MNAR：0.880 ± 0.126
  • Block-MNAR：0.908 ± 0.168
  • Seq-MNAR：0.905 ± 0.094
  • Panel-MNAR：0.791 ± 0.329（显著优于其他方法）

消融实验

零样本方法对比（表2）：

  • TabImpute+综合表现最佳：0.614 ± 0.468
  • EWF-TabPFN：0.600 ± 0.476
  • TabImpute：0.393 ± 0.487

显示了集成策略的有效性。

高缺失率场景

图4显示在MCAR模式下，随着缺失率增加，TabImpute+的优势更加明显，因为生成式模型能更好地利用上下文信息。

相关工作

缺失数据插补方法

  1. 传统方法：均值填充、线性模型、随机森林
  2. 矩阵补全：SoftImpute、USVT、近邻方法
  3. 深度学习：GAIN（GAN）、MIWAE（VAE）
  4. 集成方法：HyperImpute

表格表示学习

  1. TabPFN：表格监督学习的基础模型
  2. 后续发展：TabICL、MITRA、CausalFM等
  3. 技术特点：先验数据拟合网络（PFN）、上下文学习

结论与讨论

主要结论

  1. TabImpute实现了准确快速的零样本缺失数据插补
  2. 条目级特征化显著提升了效率和准确性
  3. 多模式训练策略增强了模型的泛化能力
  4. MissBench为插补方法评估提供了综合基准

局限性

  1. 可扩展性：由于条目级特征化，注意力复杂度在行维度上再次平方
  2. CPU性能：在CPU上运行较慢，类似TabPFN
  3. 数据类型：当前仅支持数值数据，不支持分类数据
  4. 架构限制：继承了TabPFN的二次时间复杂度限制

未来方向

  1. 探索更复杂的缺失模式和数据生成过程
  2. 增强方法以支持分类数据
  3. 扩展评估到因果推理设置
  4. 改进架构以扩展到更大数据集
  5. 利用方法进行多重插补

深度评价

优点

  1. 创新性强：条目级特征化是一个巧妙的创新，实现了显著的性能提升
  2. 实验充分：MissBench提供了迄今为止最综合的插补方法评估
  3. 实用价值高：零样本特性使得方法易于部署和使用
  4. 理论基础扎实：基于成熟的PFN框架和贝叶斯推理理论

不足

  1. 可扩展性问题：二次复杂度限制了在大规模数据上的应用
  2. 数据类型限制：仅支持数值数据限制了实际应用范围
  3. 合成数据依赖：完全依赖合成数据训练可能影响在某些真实场景下的性能
  4. 计算资源需求：训练需要大量GPU资源

影响力

  1. 学术贡献：为表格缺失数据插补领域提供了新的研究方向
  2. 实践价值：零样本特性使其在工业应用中具有很大潜力
  3. 基准贡献：MissBench将成为该领域的重要评估标准
  4. 可复现性：作者承诺开源代码和权重

适用场景

  1. 中小规模表格数据：特别适合行列数在合理范围内的数据
  2. 多领域应用：由于零样本特性，适合跨领域使用
  3. 快速原型开发：无需调参的特性适合快速验证和部署
  4. 研究和基准测试：MissBench适合作为新方法的评估标准

参考文献

本文主要基于以下重要工作：

  1. Hollmann et al. (2023, 2025) - TabPFN系列工作
  2. Müller et al. (2022) - Prior-data Fitted Networks理论基础
  3. Jarrett et al. (2022) - HyperImpute集成插补方法
  4. Rubin (1976) - 缺失数据理论基础

总结：TabImpute是一项高质量的研究工作，在技术创新、实验设计和实用价值方面都有显著贡献。尽管存在可扩展性等限制，但其零样本插补能力和优异性能使其成为该领域的重要进展。