Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Feature/fast tree partial dependence plots - #34941

Draft
ron-wettenstein wants to merge 12 commits into
scikit-learn:mainfrom
ron-wettenstein:feature/fast_tree_pdp
Draft

Feature/fast tree partial dependence plots#34941
ron-wettenstein wants to merge 12 commits into
scikit-learn:mainfrom
ron-wettenstein:feature/fast_tree_pdp

Conversation

@ron-wettenstein

@ron-wettenstein ron-wettenstein commented Sep 12, 2026

Copy link
Copy Markdown

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_accurate and recursion are timed best-of-3 after a warmup. brute is measured twice — sequential (n_jobs=1,
which is the scikit-learn default, since RandomForestRegressor ships with
n_jobs=None and _partition_estimators resolves 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)

depth grid tree_accurate recursion brute seq brute par brute seq / tree_accurate
8 5 430.5 3.1 315.6 276.2 0.7x
8 20 431.9 2.7 1258.7 1126.7 2.9x
8 100 431.9 4.0 6510.3 5515.8 15.1x
16 5 1736.1 14.3 602.5 279.7 0.3x
16 20 1746.5 27.7 2512.5 1118.4 1.4x
16 100 1765.4 63.3 12418.4 5664.8 7.0x
32 5 2324.0 24.6 700.2 281.6 0.3x
32 20 2300.6 55.6 2905.0 1097.6 1.3x
32 100 2361.9 133.6 15785.6 5704.3 6.7x

Accuracy vs brute (max absolute difference over the grid)

depth grid tree_accurate recursion
8 5 1.15e-14 7.98e-02
8 20 1.42e-14 8.66e-02
8 100 1.24e-14 8.39e-02
16 5 3.06e-14 2.77e-01
16 20 3.82e-14 2.77e-01
16 100 5.33e-14 2.77e-01
32 5 3.62e-14 2.82e-01
32 20 5.68e-14 2.82e-01
32 100 6.11e-14 2.82e-01

Notes:

  1. The method supports DecisionTree (Regressor and Classifier), ExtraTree (Regressor and Classifier), RandomForest (Regressor and Classifier). I can add GradientBoosting and HistGradientBoosting easily if needed.
  2. The approach works with multi-output regression and multi-class classification.
  3. response_method='decision_function' (or 'auto' when it directs to 'decision_function'). The approach does not support response_method='predict_proba'. I have an idea for supporting predict_proba, but it would require a completely different algorithm.
  4. I can easily support sample_weights.
  5. The approach currently compute only PDP of a single feature. Joint PDP (tuples of 2 features in the 'features' input) are not supported. I believe the approach can easily be generalized to support them.
  6. I did not touched the method="auto" routing yet. We can have it direct to "tree_accurate" instead of "recursion" when possible.

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!

@github-actions

Copy link
Copy Markdown

Thank you for opening your first pull request to scikit-learn! 🎉

To help get your contribution reviewed, please make sure that:

  • You have filled out the pull request template.

  • The pull request addresses an existing issue that is ready for contribution (e.g. not tagged as 'Needs Triage', 'Needs Decision', ...). If you are proposing a new feature, please open an issue to discuss it first.

  • There are no other open pull requests already targeting the same issue.

  • You have followed the pull request checklist. In particular, linting and tests should pass.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant