[论文解读] Multiple Imputation Through XGBoost

Yongshi Deng, Thomas Lumley|arXiv (Cornell University)|Jun 3, 2021

Statistical Methods and Bayesian Inference参考文献 20被引用 6

一句话总结

本文提出 mixgb，一种基于 XGBoost、子采样和预测均值匹配的可扩展多重插补框架，用于处理具有复杂数据结构的大规模数据集。该方法在插补准确性和计算效率方面表现卓越——尤其在 GPU 加速下——在偏差和方差方面均优于传统方法（如 mice-cart 和 mice-ranger），同时保持了统计有效性。

ABSTRACT

The use of multiple imputation (MI) is becoming increasingly popular for addressing missing data. Although some conventional MI approaches have been well studied and have shown empirical validity, they have limitations when processing large datasets with complex data structures. Their imputation performances usually rely on the proper specification of imputation models, which requires expert knowledge of the inherent relations among variables. Moreover, these standard approaches tend to be computationally inefficient for medium and large datasets. In this paper, we propose a scalable MI framework mixgb, which is based on XGBoost, subsampling, and predictive mean matching. Our approach leverages the power of XGBoost, a fast implementation of gradient boosted trees, to automatically capture interactions and non-linear relations while achieving high computational efficiency. In addition, we incorporate subsampling and predictive mean matching to reduce bias and better account for appropriate imputation variability. The proposed framework is implemented in an R package mixgb. Supplementary materials for this article are available online.

研究动机与目标

• 解决传统多重插补（MI）方法在处理具有复杂非线性关系变量的大规模数据集时的局限性。
• 开发一种可扩展、自动化的 MI 框架，能够自动捕捉交互作用和非线性效应，而无需手动指定模型。
• 提升中等和大规模数据集的计算效率，使传统 MI 方法在这些场景下变得不切实际。
• 通过引入子采样和预测均值匹配（PMM），降低插补偏差，并更好地反映不确定性。
• 评估基于 XGBoost 的插补方法在偏差、方差和运行时间方面相对于成熟 MI 方法（如 mice-cart 和 mice-ranger）的性能表现。

提出的方法

• 该框架使用 XGBoost 建模缺失变量的条件分布，自动捕捉非线性和交互效应，而无需显式指定模型。
• 在训练过程中应用子采样，以减少过拟合并提高插补的变异性，从而增强不确定性的表示。
• 使用预测均值匹配（PMM）从相似的观测案例中生成插补值，以保持数据的分布特性。
• 该方法为非迭代型，但可通过 mixgb R 包实现多轮迭代，支持 CPU 和 GPU 加速。
• 按缺失值数量对变量进行排序，并按此顺序进行插补，以提高稳定性并减少误差传播。
• 该实现以 R 包形式发布于 CRAN 和 GitHub，支持 CPU 和 GPU 执行。

实验结果

研究问题

• RQ1 在复杂数据结构中，基于 XGBoost 的插补方法结合子采样和 PMM 是否能实现比标准 MI 方法更低的偏差和更优的方差估计？
• RQ2 与 mice-cart 和 mice-ranger 相比，mixgb 的计算性能如何随数据集规模增大而变化？
• RQ3 在 XGBoost 插补中引入子采样是否能提升插补变异性并减少过拟合？
• RQ4 在估计准确性和不确定性表示方面，mixgb 相较于 mice-ranger 和 mice-cart 表现如何？
• RQ5 GPU 加速对 mixgb 的运行时间有何影响，尤其是在大规模数据集上？

主要发现

• mixgb-sub（含子采样）在所有模拟场景中均实现了最低的平均偏差，优于 mice-default、mice-cart 和 mice-ranger。
• mixgb-sub 展现出最佳的组间插补方差，表明其不确定性量化能力更优。
• mixgb-gpu 版本整体运行速度最快，尤其在大规模数据集上，子采样对性能影响可忽略不计。
• 在小规模数据集上，mixgb-cpu 比 mice-ranger 慢，但在更大数据集上，mixgb-gpu 的速度显著优于 mice-ranger。
• 即使在 MAR 机制下缺失率较高时，该方法仍能保持较高的插补质量，而完整案例分析在此类情况下会失效。
• 结果表明，对于大规模复杂数据集，mixgb-sub-gpu 是最优选择，因其在准确性、低偏差和高计算效率之间实现了良好平衡。