[Paper Review] Fast TreeSHAP: Accelerating SHAP Value Computation for Trees

Jilei Yang|arXiv (Cornell University)|Sep 20, 2021

Explainable Artificial Intelligence (XAI)Computer Science29 references44 citations

TL;DR

The paper introduces Fast TreeSHAP v1 and v2 to accelerate SHAP value computation for tree-based models, achieving about 1.5x speedup with v1 and 2.5x–3x with v2, especially benefiting multi-time interpretations at modest memory costs.

ABSTRACT

SHAP (SHapley Additive exPlanation) values are one of the leading tools for interpreting machine learning models, with strong theoretical guarantees (consistency, local accuracy) and a wide availability of implementations and use cases. Even though computing SHAP values takes exponential time in general, TreeSHAP takes polynomial time on tree-based models. While the speedup is significant, TreeSHAP can still dominate the computation time of industry-level machine learning solutions on datasets with millions or more entries, causing delays in post-hoc model diagnosis and interpretation service. In this paper we present two new algorithms, Fast TreeSHAP v1 and v2, designed to improve the computational efficiency of TreeSHAP for large datasets. We empirically find that Fast TreeSHAP v1 is 1.5x faster than TreeSHAP while keeping the memory cost unchanged. Similarly, Fast TreeSHAP v2 is 2.5x faster than TreeSHAP, at the cost of a slightly higher memory usage, thanks to the pre-computation of expensive TreeSHAP steps. We also show that Fast TreeSHAP v2 is well-suited for multi-time model interpretations, resulting in as high as 3x faster explanation of newly incoming samples.

Motivation & Objective

• Motivate the need to speed up SHAP value computation on large-scale tree ensembles in industry.
• Analyze and improve the computational steps of TreeSHAP for large sample sizes.
• Propose two variants, Fast TreeSHAP v1 and v2, with theoretical and empirical speedups.
• Assess the trade-offs between speed and memory usage across one-time and multi-time interpretation scenarios.

Proposed method

• Derive and simplify the SHAP value computation for trees to identify bottlenecks in TreeSHAP.
• Develop Fast TreeSHAP v1 by constraining the subset size tracking to subsets satisfying path thresholds, reducing average work.
• Develop Fast TreeSHAP v2 by precomputing subtree contributions U_{D_k,C} for all C ⊆ D_k to trade space for time.
• Provide complexity analyses showing theoretical time/space improvements for v1 and v2.
• Implement Prep and Score phases for v2 to enable precomputation and fast scoring on incoming data.
• Evaluate performance on multiple datasets comparing against original TreeSHAP and Fast TreeSHAP v1.

Experimental results

Research questions

• RQ1 How can TreeSHAP computation be reformulated to reduce time complexity on tree ensembles?
• RQ2 What are the practical speedups of modified TreeSHAP algorithms on large datasets and varying model sizes?
• RQ3 What is the memory-time trade-off when precomputing subtree contributions for SHAP on trees?
• RQ4 In what scenarios (one-time vs multi-time interpretations) does Fast TreeSHAP v2 provide the most benefit?

Key findings

• Fast TreeSHAP v1 yields approximately 1.5x faster computation than TreeSHAP with the same memory footprint.
• Fast TreeSHAP v2 yields approximately 2.5x–3x speedup over TreeSHAP, with a modest increase in memory usage.
• Fast TreeSHAP v2 is particularly well-suited for multi-time model interpretations due to its precomputation step.
• Across datasets and model sizes, Fast TreeSHAP v2 shows larger speedups as model size grows, with prep time increasing for larger models.
• For balanced trees, Fast TreeSHAP v2 can achieve linear-time scaling improvements in the score phase, yielding up to 3x speedups in practice.
• The SHAP value accuracy remains consistent with the original TreeSHAP up to numerical precision.