FastTreeSHAP: Accelerating SHAP value computation for trees

Co-authors: Jilei Yang, Humberto Gonzalez, Parvez Ahammad

In this blog post, we introduce and announce the open sourcing of the FastTreeSHAP package, a Python package based on the paper Fast TreeSHAP: Accelerating SHAP Value Computation for Trees (presented at the NeurIPS2021 XAI4Debugging Workshop). FastTreeSHAP enables an efficient interpretation of tree-based machine learning models by computing sample-level feature importance values, built as a new implementation of the widely-used TreeSHAP algorithm in the SHAP package. Our FastTreeSHAP package implements two new algorithms: FastTreeSHAP v1 and FastTreeSHAP v2, each improving the computational efficiency of TreeSHAP by using a different efficiency approach. Our empirical benchmarking tests show that FastTreeSHAP v1 is 1.5x faster than TreeSHAP while keeping the memory cost unchanged, and FastTreeSHAP v2 is 2.5x faster than TreeSHAP, at the cost of a slightly higher memory usage. Parallel multi-core computing is fully enabled in the FastTreeSHAP package to further speed up its computation. Our FastTreeSHAP package is easy to use, with the same API as the TreeSHAP implementation in the SHAP package, with the exception of three additional arguments which are easy to tune in practice.

Background: SHAP and TreeSHAP

Predictive machine learning models are widespread in industry today. At LinkedIn, we build predictive models to improve our member experience in different member-facing products such as People You May Know (PYMK), newsfeed ranking, search, and job recommendations, as well as customer-facing products within sales and marketing. Among these models, complex models such as random forest, gradient boosted trees, and deep neural networks are being widely used due to their high prediction accuracy. As we continue to build on our Responsible AI program at LinkedIn, a key part of our work is to understand how these models work (a.k.a. model interpretation), which remains an important challenge because these models are intrinsically opaque. 

In a previous blog post, we described how we build transparent and explainable AI systems at LinkedIn, where we highlighted a few ways we've improved transparency in AI, including explainable AI for model consumers to build trust and augment decision-making (Project Intellige), and explainable AI for modelers to perform model debugging and improvement. One of the key approaches in building transparent and explainable AI systems is to understand input contributions to model output (i.e., feature reasoning), and many times, the interpretations at an individual sample level are of the most interest. A few use cases of sample-level model interpretation at LinkedIn include:

• In our business predictive models, such as customer acquisition models and customer churn models, sample-level feature reasoning is crucial for model end users (such as the sales and marketing teams) to ensure trust in prediction results, enabling them to create meaningful insights and actionable items accordingly, which eventually leads to the improvements in our key business metrics.

• In our recruiter search models, sample-level feature reasoning can help answer questions from LinkedIn customers, such as why candidate 1 ranks higher than candidate 2, or why candidate 1’s rank has changed between searches in this month and last month, to build user trust and improve user engagement. It can also help model developers debug the model to further improve its performance. While this functionality hasn’t been implemented on the LinkedIn website yet, it is in our future plan.

• In our job search models, sample-level feature reasoning is key to supporting legal and regulatory compliance objectives, and can be helpful in ensuring our job recommendation models are fair to LinkedIn members.

There exist several state-of-the-art sample-level model interpretation approaches, e.g., SHAP, LIME, and Integrated Gradient. Among them, SHAP (SHapley Additive exPlanation) calculates SHAP values, which quantify the contribution of each feature to the model prediction by incorporating concepts from game theory and local explanations. More concretely, SHAP calculates the average impact of adding a feature to the model by accounting for all possible subsets of the other features. In contrast to other approaches, SHAP has been justified as the only consistent feature attribution approach with several unique properties (local accuracy, missingness, and consistency), which agree with human intuition. Due to its solid theoretical guarantees, SHAP has become a top model interpretation approach in industry. For more technical details of SHAP, please refer to this paper.

Figure 1 shows a typical example of SHAP values of two individual samples in the public dataset Adult, where the prediction task is to determine whether a person makes over $50K a year by using features such as marital status, educational status, capital gain and capital loss, and age. The left plot shows a prediction score of 0.776 for Person A, which is much larger than the average prediction score of 0.241, indicating a high likelihood of making over $50K a year for Person A. The top driving features are ordered from top to bottom according to their absolute SHAP values, where the red bar represents a positive value and the blue bar represents a negative value. From the left plot, we can easily see that the high capital gain and the marital status (married with a civilian spouse) contribute most to Person A’s high prediction score. Similarly, in the right plot, a prediction score of 0.007 for Person B indicates a very low likelihood of making over $50K a year, which is negatively impacted mainly by this person’s marital status (single) and young age.