Skip to content

Feature/speedup quadrature tree shap - #594

Open
ron-wettenstein wants to merge 6 commits into
mmschlk:mainfrom
ron-wettenstein:feature/speedup_quadrature_tree_shap
Open

ron-wettenstein wants to merge 6 commits into
mmschlk:mainfrom
ron-wettenstein:feature/speedup_quadrature_tree_shap

Conversation

@ron-wettenstein

Copy link
Copy Markdown
Contributor

Motivation and Context

When over the QuadratureTreeSHAP. Used Claude to better understand the code. This is an HTML describing the approach and how its is different from the algorithm in XGBoost and in the paper.
https://ron-wettenstein.github.io/TreeBranchMarks/benchmarks/resources/qshap_kernel_walkthrough.html
The new mathematical formulation is brilliant, the approach is more numerically stable than what is implemented in XGBoost. You will only feel it in depths 100+ on a strict of low coverage so I am not sure if it matters in practice - but still. In the HTML I show some upper and lower bounds on u, giving additional mathematical intuition on the stability.

In the process of reading the code I found 2 easy speedups:

  1. The path_feats variable of QuadWorkspace is used only for higher orders (not for SV) but is still maintained for SV. We can check if max_order == 1 before erasing or adding features to path_feats. This small change speed up SV by ~40%
  2. Turns out the sparse subset mapping in merged_position is the bottleneck in higher orders. The search of tuples, say (2,4,7), takes a long time, even with the smart binary search. In my measurements it took around 50% of the running time (the second heavy operation is initializing the InteractionValues object in the python - another 25%-36% of the running time). To speed up the process I encoded the tuples to integers. On 50 features in the tree, the tuple (2, 4, 7) will be encoded to 2 * 50^2 + 4 * 50 + 7. Then, we can use an int64-int32 hash map instead of a binary search. This improvement speedup SII by ~50%.
    It works as long as max_order ^ features_in_the_tree can fit in int64 (which is most cases). On other cases, we fallback to the original implementation.

Public API Changes

  • No Public API changes
  • Yes, Public API changes (Details below)

How Has This Been Tested?

The tests passes and tested manually. Add a test with the merged_position_tuple fallback.


Checklist

  • The changes have been tested locally.
  • Documentation has been updated (if the public API or usage changes).
  • An entry has been added to CHANGELOG.md (if relevant for users).
  • The code follows the project's style guidelines.
  • I have considered the impact of these changes on the public API.

@ron-wettenstein

ron-wettenstein commented Aug 27, 2026

Copy link
Copy Markdown
Contributor Author

Benchmark: quadrature kernel vs main and vs XGBoost's native TreeSHAP

XGBoost models explained three ways: by XGBoost itself, by shapiq on main, and by shapiq
on this branch. All three run on the same exported boosters and the same instance, so
only the implementation differs.

Single explanation of one instance, single thread, explainer construction excluded
(unchanged by this branch). Median of 3 rounds with the two shapiq builds interleaved,
since run-to-run spread on this machine is 5-10%.

Datasets: superconductivity (OpenML 43174, 80/20 split, 17,010 × 81, regression) and
heloc (OpenML 46932, 80/20 split, 8,367 × 23, classification). Models are
XGBRegressor / XGBClassifier, tree_method="hist", random_state=0.

Shapley values (SV)

model (trees × depth) XGBoost 3.3.0 shapiq main shapiq branch branch vs main branch vs XGBoost
superconductivity 100×4 190 us 285 us 252 us 1.13× 0.75×
superconductivity 100×8 1.28 ms 1.52 ms 946 us 1.61× 1.35×
superconductivity 100×16 10.71 ms 15.13 ms 10.51 ms 1.44× 1.02×
superconductivity 100×24 22.55 ms 31.72 ms 21.90 ms 1.45× 1.03×
heloc 100×6 415 us 674 us 516 us 1.31× 0.80×

Median: 1.44× vs main, 1.02× vs XGBoost 3.3.0.

Shapley interactions (SII, order 2)

model (trees × depth) XGBoost 3.3.0 shapiq main shapiq branch branch vs main branch vs XGBoost
superconductivity 100×4 304 us 2.52 ms 2.28 ms 1.11× 0.13×
superconductivity 100×8 3.55 ms 9.44 ms 5.49 ms 1.72× 0.65×
superconductivity 100×16 45.66 ms 82.72 ms 45.65 ms 1.81× 1.00×
superconductivity 100×24 132.43 ms 199.23 ms 113.78 ms 1.75× 1.16×
heloc 100×6 859 us 1.70 ms 1.29 ms 1.33× 0.67×

Median: 1.72× vs main, 0.67× vs XGBoost 3.3.0.

Notes

  • The SV columns are directly comparable — shapiq's Shapley values agree with XGBoost's
    pred_contribs to ≤ 7e-6 (float32 routing round-off).
  • The interaction columns are not the same quantity. XGBoost's pred_interactions
    returns Lundberg's dense symmetric SHAP interaction matrix ((F+1)², main effects on the
    diagonal); shapiq returns sparse SII of order 2. Read that column as the cost of obtaining
    pairwise interactions, not as identical output. Part of shapiq's advantage at depth is that
    its output stays sparse (3,321 interactions at 100×24 vs 6,724 dense slots).
  • Values are bit-identical between main and this branch in every case.

Environment

Windows 11, AMD Zen 2 (x86-64), single thread. Python 3.12.7, numpy 2.3.3,
scikit-learn 1.7.2. Extension built with GCC 13.2 (MinGW-w64 UCRT64) at
-O3 -ffast-math -fno-finite-math-only. Part of the gain comes from branch-prediction
effects, which may not reproduce identically on other architectures.

@codecov

codecov Bot commented Aug 27, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

This branch has not been deployed

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

Labels

None yet

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

1 participant