TreeWalker: Partial Evaluation for Grouped Tree-Ensemble Inference
Organizations: Shopify
Abstract
Many prediction tasks evaluate a trained tree ensemble on groups of rows that share feature values: discrete-time survival models expand each patient into time steps, click-through-rate models score every item in a search session, and scenario analyses vary a few inputs while holding the rest fixed. Standard inference treats each row independently and ignores this input-side correlation. We present TreeWalker, which instantiates partial evaluation for grouped tree-ensemble inference: constant features are static, varying features dynamic. TreeWalker walks each tree once per group, partitions a row bitmask at varying splits, and skips empty subtrees via an unsplit shortcut. It reads standard LightGBM and XGBoost models; training is unchanged. We prove a structural work decomposition: per-tree work splits into the constant-projected subtree size , leaf writes, and a predicate-mask provisioning cost . For the trace evaluator, per-row work approaches as , giving asymptotic speedup . On Intel, TreeWalker is faster than a row-independent traversal at the reference model size (, ; for SUPPORT/FLCHAIN and empirical groups for Expedia) and -- faster at on the survival datasets; on Arm, -- and --. On a scenario-analysis benchmark it is faster in all configurations on both machines. For f64 models, outputs match treelite GTIL up to tree-ordering roundoff; for f32 models, TreeWalker's f64 leaf accumulation is closer than native f32 to an exactly rounded reference on every audited row.
Figures & tables
| Group size | Asymptote | |||||
|---|---|---|---|---|---|---|
| Dataset | Evaluator | 4 | 16 | 64 | 128 | |
| SUPPORT | trace | 0.674 | 0.432 | 0.361 | 0.351 | 0.334 |
| SUPPORT | precompute (default) | 0.507 | 0.226 | 0.154 | 0.143 | 0.125 |
| FLCHAIN | trace | 0.544 | 0.451 | 0.439 | 0.431 | 0.417 |
| FLCHAIN | precompute (default) | 0.420 | 0.213 | 0.153 | 0.140 | 0.125 |
| Dataset | Type | Groups | Rows/Group | Const Feat. | Vary Feat. | Const% |
|---|---|---|---|---|---|---|
| SUPPORT | Survival | 1,774 | Swept (1–104) | 14 | 6 | 70% |
| FLCHAIN | Survival | 1,304 | Swept (1–1,008) | 8 | 6 | 57% |
| Expedia | CTR | 10,000 | Variable (5–37) | 10 | 11 | 48% |
| Measured ratio | Predicted | Speedup (Arm) | Speedup (Intel) | ||
|---|---|---|---|---|---|
| 1 | 0.50 | 0.195 | 0.188 | ||
| 2 | 1.15 | 0.276 | 0.268 | ||
| 4 | 2.09 | 0.395 | 0.386 | ||
| 8 | 3.81 | 0.611 | 0.601 |
Appendix figures & tables13 assets
Supplementary material from the paper’s appendix.
Appendix
| Intel | Arm | |
|---|---|---|
| GCE machine type | c4-standard-32 | c4a-highmem-16 |
| CPU | Xeon Platinum 8581C | Google Axion (Neoverse V2) |
| Frequency | 2.3 GHz base, 3.1 GHz all-core turbo | 3.0 GHz |
| Cores | 16 (SMT disabled) | 16 |
| L1d / L1i (per core) | 48 KiB / 32 KiB | 64 KiB / 64 KiB |
| L2 (per core) | 2 MiB | 2 MiB |
| Package | Version | Role |
|---|---|---|
| Rust | 1.97.1 | TreeWalker + sweep_bench |
| GCC | 15.2.0 | Native libraries, tl2cgen, lleaves linking |
| LightGBM | 4.7.0 | Model training + native predict |
| XGBoost | 3.4.1 | Model training + native predict |
| treelite | 4.7.2 | Model export (binary) + GTIL reference |
| tl2cgen | 1.0.0 | Compiled predict baseline |
| Method | Description |
|---|---|
| LightGBM native | C API: PredictForMat or a loop of PredictForMatSingleRowFast ; CMake Release with -march=native |
| XGBoost native | C API: one PredictFromDense call or a one-row loop; CMake Release with -march=native |
| tl2cgen (LightGBM+XGBoost) | Compiled C trees via treelite (GCC -O3 -march=native ); its runtime or a loop of the generated predict |
| lleaves (LightGBM) | LLVM-compiled trees ( -O3 code generation for the host CPU; trees not inlined into the forest loop) |
| QuickScorer (LightGBM+XGBoost) | Bitvector-based predictor ( quickscorer-rs ), on the encoded Expedia variant and the five original-Expedia cells of LightGBM’s , model (Appendix D ) |
| TreeWalker | Partial evaluation (this work) |
| Field | Type | Description |
|---|---|---|
| value | f64 | Threshold, inline category bitset or bitset-pool offset, or leaf output |
| skip | i16 | Tree-local index of the light child ( = leaf) |
| varying_pred_id | u16 | Index into varying-predicate array |
| feature | u16 | Feature index |
| flags | u8 | Bitfield: default_left, categorical, inline_cat, heavy_is_left, walkable, vary_type (2 bits: constant, increasing, decreasing, non-monotonic) |
| cat_n_words | u8 | Bitset word count (categorical splits) |
| Term | Counter (proxy) | Source |
|---|---|---|
| constant_steps | per-tree stats (exact) | |
| varying_splits | per-tree stats (exact) | |
| group size | known (exact) | |
| precompute_row_evals partition_row_evals | per-group stats (proxy) | |
| offline replay against | one-time post-process |
| Framework | TreeWalker ( s) | QuickScorer ( s) | QS/TW (Intel) | QS/TW (Arm) | ||
|---|---|---|---|---|---|---|
| LightGBM | 50 | 2 | 5.38 | 3.44 | ||
| LightGBM | 500 | 4 | 67.13 | 61.22 | ||
| LightGBM | 1000 | 4 | 107.35 | 105.35 | ||
| XGBoost | 50 | 2 | 6.85 | 3.90 | ||
| XGBoost | 50 | 4 | 16.95 | 10.35 | ||
| XGBoost | 500 | 2 | 34.85 | 24.09 |
| Speedup (Arm) | Speedup (Intel) | Trace ratio | Predicted | Precompute ratio | ||
|---|---|---|---|---|---|---|
| 1 | 4 | 0.413 | 0.188 | 0.351 | ||
| 1 | 16 | 0.245 | 0.188 | 0.183 | ||
| 1 | 64 | 0.202 | 0.188 | 0.140 | ||
| 1 | 128 | 0.195 | 0.188 | 0.133 | ||
| 2 | 4 | 0.502 | 0.268 | 0.359 | ||
| 2 | 16 | 0.329 | 0.268 | 0.186 |
| Rows per group | Intel | Arm |
|---|---|---|
| [ , ] | [ , ] | |
| [ , ] | [ , ] | |
| [ , ] | [ , ] |
Efficient Inference
Efficient inference reduces the compute, memory, latency, or energy needed to produce accurate predictions, posterior estimates, and decisions.
See all 109 papers