Feature/speedup quadrature tree shap - #594
ron-wettenstein wants to merge 6 commits into
Conversation
Benchmark: quadrature kernel vs
|
| 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_contribsto ≤ 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
mainand 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 Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
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:
merged_positionis 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
How Has This Been Tested?
The tests passes and tested manually. Add a test with the merged_position_tuple fallback.
Checklist
CHANGELOG.md(if relevant for users).