Feature/fast tree partial dependence plots - #34941
Draft
ron-wettenstein wants to merge 12 commits into
Draft
Conversation
|
Thank you for opening your first pull request to scikit-learn! 🎉 To help get your contribution reviewed, please make sure that:
|
ron-wettenstein
marked this pull request as draft
September 12, 2026 12:50
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This is a draft PR that relates to the issue #28771 .
The partial_dependence function has two methods: "brute" that uses brute force to compute the values exactly and "recursion" that estimate these values using the tree structure.
The issue discuss the differences in the values returned by the "recursion" and "brute" methods. In fact, a recent paper I published: https://arxiv.org/abs/2605.14578, shows that the "recursion" method uses same cooperative game as Path Dependent SHAP while the values returned by the "brute" method are defined over the same game as Interventional SHAP.
This PR introduce the new "tree_accurate" method.
The method does the best of both worlds - it produce the same values as the "brute" while using the tree structure to speed up the computation. The approach is often faster than "brute". Unfortunately is slower than the "recursion" method as it considers the provided data and not just the tree structure.
The algorithm
Let f1 denote the feature whose partial dependence values we want to compute.
We start with a preprocessing traversal of the tree. This pass considers the background data X and have O(|X| * depth) complexity. (In the extreme case where all nodes split on f1, the complexity is O(|X| * K) where K is the number of nodes in the tree).
For each leaf, we count how many rows from X reach it. For inner nodes that split on f1 the row traverses both right and left. For inner nodes that split on another feature the row traverses right/left according to its split result.
Then we use this count to compute the partial dependence values. This pass consider the grid values and have O(|grid| * K) complexity.
For each grid value we traverse the tree, going both right and left on splits that does not use f1 and right/left according to its split result for splits that use f1. Each reached leaf contributes to the value: leaf_weight * count[leaf]/|X|
The algorithm is implemented in the compute_partial_dependence_tree_accurate function in _tree.pyx
Benchmarking
For benchmarking, I used the California housing dataset. I split it 80/20 with random_state=0, the trainset have 16512 rows and 8 features. I trained RandomForestRegressor models with 100 trees and depths 8/16/32 on this trainset and compute the partial dependence values of MedInc (the dominant feature) using difference grid resolutions and the entire trainset as reference data.
tree_accurateandrecursionare timed best-of-3 after a warmup.bruteis measured twice — sequential (n_jobs=1,which is the scikit-learn default, since
RandomForestRegressorships withn_jobs=Noneand_partition_estimatorsresolves that to 1) and parallel(
n_jobs=-1, spreading the 100 trees across cores).I run it on my laptop: AMD Ryzen 7 5700U (8 physical / 16 logical cores, 1.8 GHz base), 37.8 GB RAM. Python 3.12.7, NumPy 2.4.4, SciPy 1.17.1, joblib 1.5.3, scikit-learn 1.10.dev0
Timings (ms)
Accuracy vs
brute(max absolute difference over the grid)Notes:
About myself: I'm a PhD student researching decision trees explainability. In my PhD I develop faster algorithms for computing explainability methods such as Shapley values, Banzhaf values and partial dependence plots on decision trees.
I recently worked on the Woodelf algorithm (see https://github.com/ron-wettenstein/woodelf) that computes interventional Shapley/Banzhaf values and Shapley/Banzhaf interaction values, improving the state-of-the-art complexity. Woodelf can also compute partial dependence plots, as discussed in the paper: https://arxiv.org/abs/2605.14578, but I think the approach is not a good fit to scikit-learn, as scikit-learn's partial_dependence function accept one feature at a time and Woodelf prefer computing the partial dependence values of all features in one go.
I used claude code (Opus and Fable) for implementation and testing. I read, understood and improved all the code in this PR.
I will be happy to hear your thoughts on this!