[论文解读] GPUTreeShap: Massively Parallel Exact Calculation of SHAP Scores for Tree Ensembles

Rory Mitchell, Eibe Frank+1 more|arXiv (Cornell University)|Oct 27, 2020

Machine Learning and Data Classification被引用 8

一句话总结

GPUTreeShap 提出了一种针对树模型集成的 GPU 优化、大规模并行重构的 TreeShap 算法，用于精确计算 SHAP 值。通过将递归子问题重构成基于 bin-packing 调度的 SIMT 任务，其在 SHAP 值计算上相比多核 CPU 最快提升 19 倍，在交互值计算上最快提升 340 倍，实现了大规模场景下的实时可解释性。

ABSTRACT

SHAP (SHapley Additive exPlanation) values provide a game theoretic interpretation of the predictions of machine learning models based on Shapley values. While exact calculation of SHAP values is computationally intractable in general, a recursive polynomial-time algorithm called TreeShap is available for decision tree models. However, despite its polynomial time complexity, TreeShap can become a significant bottleneck in practical machine learning pipelines when applied to large decision tree ensembles. Unfortunately, the complicated TreeShap algorithm is difficult to map to hardware accelerators such as GPUs. In this work, we present GPUTreeShap, a reformulated TreeShap algorithm suitable for massively parallel computation on graphics processing units. Our approach first preprocesses each decision tree to isolate variable sized sub-problems from the original recursive algorithm, then solves a bin packing problem, and finally maps sub-problems to single-instruction, multiple-thread (SIMT) tasks for parallel execution with specialised hardware instructions. With a single NVIDIA Tesla V100-32 GPU, we achieve speedups of up to 19x for SHAP values, and speedups of up to 340x for SHAP interaction values, over a state-of-the-art multi-core CPU implementation executed on two 20-core Xeon E5-2698 v4 2.2 GHz CPUs. We also experiment with multi-GPU computing using eight V100 GPUs, demonstrating throughput of 1.2M rows per second -- equivalent CPU-based performance is estimated to require 6850 CPU cores.

研究动机与目标

• 解决大规模树模型集成中 SHAP 值计算带来的计算瓶颈，其耗时可能超过模型训练时间。
• 克服由于不规则递归和数据依赖性，将复杂 TreeShap 算法映射到 GPU 架构的挑战。
• 通过利用 GPU 的大规模并行性，实现实时 ML 工作流中高吞吐、低延迟的 SHAP 计算。
• 降低 SHAP 交互值计算的时间成本，其时间复杂度随特征数量呈平方级增长，尤其在高维数据中开销极大。
• 为 XGBoost 提供一个可投入生产的、开源的 GPU 后端，实现跨多个 GPU 的高效扩展。

提出的方法

• 将递归的 TreeShap 算法重构为每条树路径独立的可变大小子问题，以支持数据并行执行。
• 应用装箱启发式算法将子问题分组为线程束（warps），以实现高效的 GPU 线程块调度并最大化占用率。
• 使用 CUDA 内核将每个子问题映射为单指令多线程（SIMT）任务，并利用线程束级原语支持动态规划。
• 利用训练数据中的覆盖权重，在路径评估期间估计缺失特征的条件期望。
• 通过省略无关特征对，优化 SHAP 交互值的算法复杂度，将时间复杂度从 O(TLD²M) 降低至 O(TLD³)。
• 将 GPUTreeShap 集成为 XGBoost 的后端，实现无需重新训练模型的无缝 GPU 加速 SHAP 推理。

实验结果

研究问题

• RQ1 递归的 TreeShap 算法能否被重构为在大规模并行 GPU 架构上实现高效可扩展执行？
• RQ2 GPU 级并行性和线程束调度在多棵大型树模型的 SHAP 值计算中，能在多大程度上提升吞吐量？
• RQ3 GPUTreeShap 的性能如何随模型规模、数据集大小和 GPU 数量的增加而扩展？
• RQ4 对于显著比标准 SHAP 值更昂贵的 SHAP 交互值，可实现多大的性能提升？
• RQ5 能否通过预处理和特征剪枝，在不损失精度的前提下降低 SHAP 交互计算的算法复杂度？

主要发现

• 在单张 NVIDIA Tesla V100-32 GPU 上，GPUTreeShap 相比 40 核 CPU 集群，标准 SHAP 值计算最快提升 19 倍，交互值计算最快提升 340 倍。
• 吞吐量随 GPU 数量线性增长；八张 V100 GPU 可实现每秒 120 万行的处理速度，相当于约 6,850 个 CPU 核心的性能。
• 对于高维数据集如 fashion_mnist-large，使用单张 GPU 可将 SHAP 交互值计算时间从约 6 小时缩短至约 1 分钟。
• 通过特征子集剪枝，SHAP 交互值的算法复杂度从 O(TLD²M) 降低至 O(TLD³)，显著提升了大规模特征集上的性能。
• 在大批次输入下，GPUTreeShap 展现出更高的吞吐量，当测试样本超过 180 行后，其性能超越 CPU，得益于更高的并行度和更低的每行延迟。
• 该实现已具备生产就绪能力，并已集成至 XGBoost，为梯度提升树模型提供了端到端的 GPU 加速可解释性支持。