Skip to content

About

No description, website, or topics provided.

Resources

Code of conduct

Security policy

Stars

2 stars

Watchers

0 watching

Forks

Repository files navigation

TreeWalker

Code and data for TreeWalker: Partial Evaluation for Grouped Tree-Ensemble Inference (Durmus Karatay and Richard Newman, NeurIPS 2026).

TreeWalker evaluates a trained tree ensemble on groups of rows that share most feature values. It walks each tree once per group, partitions a row bitmask at splits on varying features, and skips empty subtrees. This repository contains the engine, the benchmark harness, the treewalker-exp package that turns public datasets into models, runs the benchmark suites and computes the paper's figures, tables and numbers, the raw measurements of the final run, and the cloud setup used for them.

The measurements here are the 2.0 engine's final run (October 2026). The paper as submitted, its engine, its result CSVs and the scripts that read them are preserved under the neurips2026-submission tag; the camera-ready paper's code, measurements and figure code are under neurips2026-final.

This repository serves two purposes: as a source for the library and as a permanent archive of the code at time of submission. The latter can always be found under the neurips2026-submission tag.

Layout

Path Contents
src/ inference library and treelite model parser
experiments/benchmarks/ separate unpublished crate: benchmark harness, sweep_bench, artifact correctness tests
experiments/treewalker_exp/ treewalker-exp: dataset preparation, workloads, training, compiled baselines
experiments/grids.toml the benchmark suites, resolved by treewalker-exp into execution manifests
experiments/data/runs/ the final run's raw measurements: packed runs, Parquet in git-lfs
experiments/figures/ figures, tables and numbers treewalker-exp generates (git-ignored)
infra/ Terraform and VM startup script for the two GCE benchmark machines

Building and packaging the library

The root crate, treewalker-gbdt, contains the inference library. Rust imports use treewalker_gbdt, for example use treewalker_gbdt::Forest;. These commands build, test, and package it without the benchmark dependencies:

cargo build
cargo test --release
cargo package

cargo test also runs the research API's tests: the crate's dev-dependency on itself turns the research feature on. Version 2.0 changed the API; the changelog maps every 1.x call to its replacement.

The package includes the Rust library sources, loading documentation, README, license, and small independent compatibility fixtures with their optional Python generators. Benchmark sources, result data, experiment scripts, Terraform files, and the repository's native CPU build settings are excluded.

The benchmark crate, treewalker-bench, is a workspace member with its own manifest and the shared Cargo.lock. Build it explicitly; --target-dir target preserves the binary path used by the experiment scripts:

cargo build --manifest-path experiments/benchmarks/Cargo.toml --target-dir target \
  --release --features external-bench

The external-bench feature enables the C baselines and QuickScorer; quickscorer-bench enables QuickScorer alone, and pmu the Linux hardware counters. cargo bench runs small divan benchmarks of the library on the committed fixtures.

Development checks

Every commit passes cargo fmt --check, cargo clippy --workspace --all-targets -- -D warnings and cargo test. The Linux-only code (the timestamp counter's calibration, perf_event) is checked by cross clippy. Parquet's zstd is C, so without a Linux cross toolchain the check compiles it with clang and the macOS SDK headers; nothing builds or links for Linux:

SDK=$(xcrun --show-sdk-path)
export CC_x86_64_unknown_linux_gnu=clang \
  CFLAGS_x86_64_unknown_linux_gnu="--target=x86_64-unknown-linux-gnu -isystem $SDK/usr/include" \
  CC_aarch64_unknown_linux_gnu=clang \
  CFLAGS_aarch64_unknown_linux_gnu="--target=aarch64-unknown-linux-gnu -isystem $SDK/usr/include -D__arm64__"
for t in "aarch64-apple-darwin apple-m4" "x86_64-unknown-linux-gnu emeraldrapids" \
         "x86_64-unknown-linux-gnu x86-64" "aarch64-unknown-linux-gnu neoverse-v2"; do
  set -- $t
  for f in "" "--features treewalker-bench/external-bench,treewalker-bench/pmu"; do
    RUSTFLAGS="-C target-cpu=$2" cargo clippy --target $1 --workspace --all-targets $f -- -D warnings
  done
done

Using the library

Load a model once and predict through a Predictor, one per worker thread. Forest is a cheap handle to the immutable model that can be cloned and shared; a predictor owns the scratch space for the widest group, so its calls never allocate.

use treewalker_gbdt::{Forest, LoadError};

fn main() -> Result<(), LoadError> {
    let forest = Forest::load("model.bin", "walker_config.json")?;
    let mut predictor = forest.predictor();
    // Two entities with 3 and 2 rows, row-major, in the model's trained column order.
    let data = vec![0.0; 5 * forest.config().n_features()];
    let mut out = vec![0.0; 5];
    predictor.predict_groups(&data, &[0, 3, 5], &mut out);
    Ok(())
}

predict_group predicts one group and predict_fixed consecutive groups of one width. Forest::from_reader and from_bytes accept an explicit ModelFormat and a WalkerConfig; build one in code with WalkerConfig::builder(n_features).max_group_width(w).varying([..]).build()?. The builder requires the varying features, so a forgotten list is an error rather than every feature silently treated as constant.

The supported subset includes scalar regression, ranking and binary classification, identity/sigmoid, sum/average aggregation, and scalar base scores. Binary Treelite v4 is recommended; native XGBoost JSON must first be converted through Treelite. Inputs have up to 64 features; groups can have any number of rows. See Treelite loading for export examples, the exact output and precision policy, errors, limits, and the caller's grouping contract.

Datasets

None of the datasets are redistributed here. uv run treewalker-exp fetch downloads the public ones into experiments/data/raw/ and checks each against a pinned SHA-256; prepare fetches them on demand.

Dataset Source Fetched by
SUPPORT DeepSurv repository (support_train_test.h5) treewalker-exp fetch
FLCHAIN R survival package via Rdatasets treewalker-exp fetch
Expedia (ICDM 2013) Kaggle competition expedia-personalized-sort treewalker-exp fetch-expedia, from the data.zip you download
UCI Default of Credit Card Clients OpenML 42477, CC BY 4.0, DOI 10.24432/C55S3H treewalker-exp fetch

The Expedia data may not be redistributed. Accept the competition rules on Kaggle, download data.zip (for example kaggle competitions download -c expedia-personalized-sort -f data.zip), then run treewalker-exp fetch-expedia --train-csv data.zip; extracting the zip needs Info-ZIP unzip, because it uses Deflate64. It writes experiments/data/expedia.parquet and checks its content fingerprint against the file the paper used.

The factorial suite also prepares two kinds of derived models (experiments/grids.toml), neither of which changes an existing model or cell:

  • expedia-filled, the Expedia split with every missing value encoded before training as its column's training-split minimum minus max(1, |minimum|), the same constant in train and test (test values below it are kept). The constant stays below the minimum after the f32 conversion, so QuickScorer, which has no missing-value handling, can take these models. That does not make it compute the same function: it compares thresholds in f32, and in the final run 5 of the 16 cells' QuickScorer outputs failed validation against LightGBM's f64 thresholds. Only cells whose QuickScorer outputs validate are timed. Sessions workload, T in {50, 500, 1000, 2000}, L in {2, 4}, with Expedia's sessions, order and varying features.
  • Seed replicates, _r1 to _r4 after a model's name. Training is deterministic, so each replicate trains with the released parameters on a seeded sample of exactly 80% of the training split's entities (patients, sessions, applicants) and is timed on the released cell's data. At three anchors: support/nt500_md4_h16 (panel), credit/nt500_md4 (what-if k4 G16, on the features the released model splits on most) and expedia/nt500_md8 (sessions). A model's model.json records the sample (replicate) or the encoding (missing_values).

Released results

experiments/data/runs/ holds the camera-ready campaign (2026-10-11 to 13, exact sums off) and the final run's factorial and ablation runs (2026-10-09 and 10, exact sums on), one directory per run ID, packed with treewalker-exp pack: run.json (configuration, host, timer, sentinel and the order cells ran in), cells.json (each cell's manifest: its sample, validation, stopping, probes and XGBoost's processes), the samples, groups, counters, hw and stages tables, timer.parquet, and the sentinel's cells under sentinel/. The tables are Parquet in git-lfs: run git lfs pull after cloning. stages (per cell stages.parquet; packed as stages.parquet and sentinel/stages.parquet) holds TreeWalker's tree sum, raw margin and output on each row of the cell's stage oracle (columns suite, cell, group, row, tree_sum, raw_margin, output), written untimed before timing; output schema 6, absent from older runs.

Run Campaign Contents
camera-ready-{x86_64,aarch64} camera-ready the paper's 712 cells: survival panels (SUPPORT and FLCHAIN, horizons 1-128; FLCHAIN also 256-1,024, as whole groups and as 128-row chunks), Expedia sessions, cohorts of 4-32 rows and filled missing values
camera-ready-ablation-{1,2}-{x86_64,aarch64} camera-ready the 36 ablation cells, run before and after camera-ready
scenario-v1-{x86_64,aarch64} camera-ready E1: the 16 released credit scenario cells
acceptance-{x86_64,aarch64} camera-ready 17 acceptance cells
layout-{default,align64}-{x86_64,aarch64} camera-ready the acceptance cells with default and 64-byte-aligned functions
RUN-retime-rN-{x86_64,aarch64} camera-ready re-time attempts: the cells and modes of RUN whose interval was wider than ±2%, timed again with N rounds on the same VM
factorial-{x86_64,aarch64} final 1,134 cells: survival panels (SUPPORT and FLCHAIN, horizons 1-128; FLCHAIN also 256-1,024), Expedia sessions, cohorts of 4-32 rows and filled missing values, the credit and reference what-ifs, and seed replicates
ablation-{x86_64,aarch64} final 36 cells with every switch and their combinations, in serving and batch mode

experiments/data/retime/ holds each re-time run's manifest and script. The final run's acceptance and layout runs, which share their IDs with the camera-ready ones, are in the history at 2562892.

experiments/data/precision/{intel,arm}/ holds App. A.3's measurements, taken on the campaign's machine types (2026-10-13, the precision diagnostic): precision_order.json, sweep_bench precision-order over every camera-ready LightGBM cell (tree ordering on against off with exact sums off, and both against prepare's GTIL reference); audit_f32.md, audit-f32 with native XGBoost's margins from the library the campaign timed (xgboost_library.txt records its hashes), beside the camera-ready runs' TreeWalker stages; and both commands' logs.

Reproducing the tables and figures

This needs Python 3.14 (.python-version), uv and git-lfs, not the benchmark machines. From the repository root:

git lfs pull
uv sync
uv run treewalker-exp figures   # Figures 2-9: PNG and PDF, and the heatmap as TikZ
uv run treewalker-exp tables    # Table 1 (TeX), Tables 3 and 10, App. A.5 and D, ablation runs (Markdown)
uv run treewalker-exp numbers   # every in-text number, with its paper location
uv run treewalker-exp summarize experiments/data/runs/factorial-x86_64   # one run, cell by cell

Outputs go to experiments/figures/. figures, tables and numbers read a campaign's runs (--campaign): final, the run above, or camera-ready (below); without the option, the camera-ready runs where they exist. The f32 audit also needs prepared models and cells under experiments/artifacts/; treewalker-exp audit-f32 --list-missing lists the files it reads.

Paper item Command Output
Figure 2 (heatmap) figures speedup_heatmap.{pdf,png,tex}
Figures 3, 4, 6-9 figures horizon_amortization, method_comparison, cross_platform, ablation_waterfall, ablation_nodes, depth_sensitivity, ntrees_scaling
Table 1, Figure 5 tables, figures decomposition_table.tex, asymptotic_convergence
Tables 3 and 10 tables scenario_credit_{intel,arm}.md
App. A.5 (tl2cgen's f64 aliasing defect) tables tl2cgen_defect_{intel,arm}.md
App. D (QuickScorer on expedia-filled, replacing the withdrawn Table 9) tables quickscorer_filled_{intel,arm}.md
§5.7 ablation effects per run and their between-run range tables, figures ablation_effects.md, ablation_waterfall
Table 11, in-text numbers (§4-§7, App. A.6, D, F) numbers paper_numbers.txt
App. A.3 (f32 audit) audit-f32 audit_f32.md

A method's latency in a cell is v1's estimand, completed for rounds: the median over blocks (one round's pass over one batch) of each block's median time per group, by v1's nearest-rank convention, from process 0; p5-p95 over blocks is descriptive dispersion. A speedup is the ratio of two latencies, with a 95% bootstrap interval that resamples entities and rounds as crossed clusters, pairs the two methods and recomputes the nested medians. Sampled cells give every group the same inclusion probability, which the analysis checks, so the medians are unweighted. The row-weighted ratio of mean time per row and its interval are kept as a sensitivity column (speedup_rw), and XGBoost native's faster mode across its extra processes as a sensitivity row. Every cell runs in a process of its own, so no method is corrected for drift. numbers keeps the v1 scripts' items, and its module docstring says where an item changed.

Rerunning the benchmarks

The final run ran on two Google Compute Engine VMs in us-east4-a: c4-standard-32 (Intel Xeon Platinum 8581C, at its all-core turbo) and c4a-highmem-16 (Google Axion, Neoverse V2), Ubuntu 26.04, SMT disabled, ASLR off, hardware counters at the PMU's STANDARD level, pinned to one core with taskset -c 0. LightGBM 4.7.0 and XGBoost 3.4.1, the versions uv.lock installs, were built from source for the host, and TreeWalker with -C target-cpu=native and Rust 1.97.1 (rust-toolchain.toml). Each run.json records the rest (host, system_info). infra/scripts/startup.sh is the full machine recipe. The paper's measurements (the neurips2026-submission tag) ran on the same machine types in us-central1-a.

On GCE with Terraform

cd infra
terraform init
terraform apply -var project=YOUR_PROJECT -var git_ref=$(git rev-parse HEAD) -var 'suites=["acceptance"]' -var layout_check=true
terraform apply -var project=YOUR_PROJECT -var git_ref=$(git rev-parse HEAD)      # factorial and ablation

With -var cache_bucket=NAME, an existing bucket keeps the prepared models and each machine type's compiled baselines across deployments, so a rerun retrains and recompiles only what changed. Prep and compile-baselines reuse a cached file only when its recorded key and hash match. Create the bucket once, outside Terraform, so terraform destroy leaves it: gcloud storage buckets create gs://NAME --location US --uniform-bucket-level-access.

The VMs run with the GCE PMU at STANDARD and turbo_mode = "ALL_CORE_MAX". They boot a pinned image, and every apt operation uses one Ubuntu archive snapshot (apt_snapshot), so the compiler and tools, and with them the cached compiled baselines, stay the same across deployments. Suites with Expedia cells need experiments/data/expedia.parquet from treewalker-exp fetch-expedia; Terraform uploads it for the trainer VM, which checks its fingerprint before training. Terraform uploads git archive of var.git_ref to a new bucket. git_ref is a full 40-character commit SHA, which Terraform checks: the archive, the VMs and the model hand-over are named after it, so a moving branch would keep the old source. It has no default: it must contain experiments/ and treewalker-exp, which the neurips2026-submission tag predates. The Intel VM trains the models and shares them, so both machines evaluate identical models, checked by hash. Results land in the bucket under results/.

By hand on a prepared machine

uv sync --group baselines
uv run --group baselines treewalker-exp build-native            # LightGBM, XGBoost from pinned sources
uv run --group baselines treewalker-exp fetch-expedia --train-csv PATH/data.zip
uv run --group baselines treewalker-exp prepare --suite acceptance
uv run --group baselines treewalker-exp compile-baselines --suite acceptance   # tl2cgen, lleaves, QuickScorer XML
uv run --group baselines treewalker-exp preflight --suite acceptance
uv run --group baselines treewalker-exp run --suite acceptance --run-id acceptance
uv run --group baselines treewalker-exp pack experiments/data/runs/acceptance
uv run --group baselines treewalker-exp summarize experiments/data/runs/acceptance
uv run --group baselines treewalker-exp check-run experiments/data/runs/acceptance

Then the same with --suite factorial and --suite ablation. run resolves the suite into an execution manifest, builds sweep_bench and runs it pinned to core 0 with taskset where it exists; a rerun with the same run ID reuses the finished cells whose manifests match. The runner writes run.json and, per cell, manifest.json and the samples, groups, counters, hw and stages Parquet tables. pack then writes one file per table and one cells.json for the run (the sentinel's cells under sentinel/), checks every table against its cells, and with --remove deletes only the finished cell directories it packed and verified, keeping unfinished cells, failed sentinel attempts and the run's failed/; from output schema 6 it requires stages.parquet. summarize, validation-readout and budget read packed runs. run --set KEY=VALUE overrides a grids.toml [run] setting, after the suite's own [suites.NAME.run] settings (the validation suite's).

The protocol, per cell (settings in grids.toml [run]):

  • Each suite declares its modes: serving (one call per group) everywhere, batch (calls of up to 1,024 rows) in the ablation suite as well.
  • Methods run one at a time. The timed groups are split into 12 batches; a round gives each method its own phase, a warm-up on the batch it times last and then a timed pass over every batch, the same groups in the same order for every method. Method order alternates across rounds, forward then reversed (A B B A). A mode times at least 3 rounds, then stops at the 1% precision target or its time budget.
  • A cell whose pool exceeds max_rows_per_round rows times a seeded sample of its groups, the same for every method: every group has the same inclusion probability, the cap's share of the pool's rows (the cap is the expected rows), raised to keep at least min_groups_per_round groups, stratified by group size. The manifest records the seed, the fraction, the drawn rows and their overshoot, where the minimum binds, the strata with their weights and a hash of the group indices.
  • Baselines get their faster interface per cell and mode: in its first phase, each method's candidates (its multi-row call, or a loop over its single-row call) run on the warm-up batch, and the faster is timed. One-row serving calls use the single-row interface. The manifest records each probe.
  • QuickScorer is excluded up front from LightGBM models with categorical splits, and from cells with missing values in a feature the model may route otherwise than left (a conservative rule: QuickScorer sends every missing value left). Every method receives f64 row-major input and converts it inside its timed call; QuickScorer's f64-to-f32 copy is part of its samples, and is also timed alone, right after its timed passes, so summarize reports its share.
  • On x86_64, each XGBoost cell also times XGBoost alone in 6 extra one-round processes (process in the samples), because its AVX2 block walk depends on a heap buffer's alignment. XGBoost's main number is process 0's, timed and paired like every method. As a sensitivity row, summarize tells the modes apart inside each batch by instructions per row and reports the faster one over the batches holding both, with the expected time across processes beside it; the extra processes validate only their first 32 groups.
  • Each cell, and each sentinel, runs in a new process of the runner (fresh_process). In one long-lived process, LightGBM native ran up to 6.9% slower by the end of an earlier factorial, with the same instructions: a model loaded into a glibc heap that hundreds of models had churned. A new process, or jemalloc, showed none.
  • A sentinel, SUPPORT's acceptance cell, runs at the start of each run invocation (a warm-up, then a measure) and after every 25 measured cells, under sentinel/, its ticks and rows in run.json. It measures the machine; summarize flags drift beyond 2% from the run's first measured sentinel and shows the first-load effect, and numbers reports each method's largest drift.

compile-baselines runs as many compiler processes at once as there are CPUs (--jobs, at least 4), largest models first, tl2cgen at 2 threads per model and lleaves at its recipe's 4 chunks. The chunk count shapes lleaves's library, so it is part of the recipe: on the 32-vCPU Intel VM the acceptance and shakedown runs compiled lleaves in 8 chunks, the final run in 4. Running the largest models first puts the biggest compiles side by side; their peak memory is checked on the validation deployment. The artifact tests run on whatever was prepared: cargo test --manifest-path experiments/benchmarks/Cargo.toml --release --features research.

tl2cgen comes from the locked baselines dependency group. It has no aarch64 Linux wheel, so there it builds from source with CXX=$PWD/infra/scripts/cxx-cstdint. lleaves runs in experiments/compile.py's own script lock and compiles through the LLVM that llvmlite bundles; the system C compiler links it. Every uv run passes --group baselines, because uv run removes packages outside the groups it syncs.

uv.lock, experiments/compile.py.lock and tests/fixtures/import/generate.py.lock resolve against PyPI through the pypi index that the project and both scripts declare, which takes priority over any index in a user's uv configuration. Relock with uv lock, uv lock --script experiments/compile.py and uv lock --script tests/fixtures/import/generate.py; uv lock --check (with --script for a script) verifies a lock without changing it.

The camera-ready rerun

The camera-ready reruns the paper's cells with the algorithm as reviewed: every TreeWalker forest loads with exact sums off (exact_sums = false), so each row's leaves add in f64 in tree order, as Algorithms 1 and 2 describe, and production predict equals predict_full_walk bit for bit. The library's default stays on (LoadOptions::exact_sums, true by default: the leaves add as exact fixed-point integers rounded once; see docs/treelite-loading.md). In the final run's ablation suites (ablation-{x86_64,aarch64}, measured 2026-10-09/10, numbers item README), turning exact sums off (disable_exact_sums, research build, serving, T=500, L=8) took 1.26-1.29x the default's latency at G=128 on Intel and 1.59-1.60x on Arm, 1.25-1.28x and 1.40-1.45x at G=64, 1.10-1.15x and 1.13-1.15x at G=16, 1.02-1.06x and 0.99-1.02x at G=4, and 0.98-0.99x and 0.96-0.97x on Expedia's sessions, where exact sums are the slower. Each mode times at least 4 rounds (min_rounds = 4), two full forward-and-reversed order cycles. The suites, in grids.toml, each set both in their own run table, so factorial and ablation keep their defaults:

  • camera-ready: survival panels at horizons 1-128 and FLCHAIN's 256-1,024 (128-row chunks beside whole groups), Expedia's sessions, the fixed-model cohorts of 4-32 rows and expedia-filled.
  • camera-ready-ablation: the ablation's anchors and switches without the disable_exact_sums variant, run twice as two runs.
  • scenario-v1: E1, the released credit scenario cells, with the trace evaluator for Table 3's d_v columns.

infra/camera-ready.tfvars holds the campaign: acceptance, the ablation, the camera-ready suite, the ablation again, scenario-v1 and the layout check, every run with --set exact_sums=false:

cd infra
terraform apply -var project=YOUR_PROJECT -var git_ref=$(git rev-parse HEAD) -var 'suites=["acceptance"]' -var layout_check=true -var 'run_settings=["exact_sums=false"]'
terraform apply -var project=YOUR_PROJECT -var git_ref=$(git rev-parse HEAD) -var-file=camera-ready.tfvars

The runs land as camera-ready-ARCH, camera-ready-ablation-1-ARCH, camera-ready-ablation-2-ARCH and scenario-v1-ARCH. By hand, the same suites run with treewalker-exp run --suite NAME --run-id NAME-ARCH --set exact_sums=false (the second ablation run as camera-ready-ablation-2-ARCH). Then, per run:

uv run treewalker-exp pack experiments/data/runs/RUN
uv run treewalker-exp check-run experiments/data/runs/RUN ...   # gates and warnings
uv run treewalker-exp retime camera-ready                       # wide intervals, twice the rounds (below)
uv run treewalker-exp figures --campaign camera-ready           # and tables, numbers
uv run treewalker-exp audit-f32 --suite camera-ready            # f32 margins against the oracle
cargo run --release --manifest-path experiments/benchmarks/Cargo.toml --features external-bench \
  --bin sweep_bench -- precision-order --manifest MANIFEST --output order.json   # offline: ordering on vs off

audit-f32 --suite camera-ready reads TreeWalker's f32 margins from each run's packed stages.parquet and compares them, and native XGBoost's, with the same f64 oracle. sweep_bench precision-order runs offline over prepared LightGBM cells (cell directories, or --manifest with --cells and --artifacts-dir): with exact sums off it compares tree ordering on against off on the tree sums and raw margins, and both orderings' outputs with the cell's GTIL reference (experiments/benchmarks/src/precision.rs, docs/treelite-loading.md).

check-run exits non-zero when a validity gate fails, and fails closed: the run's tables missing, a failed cell or sentinel attempt that no later attempt of the same cell (in the order of run.json's invocations, each failure in failed/ recording its invocation), or of the sentinel's role, resolved (resolved failures stay in the data and are listed), a cell whose stage oracle did not pass, a required method not timed, a timed interface without a passing check on the timed rows, process 0's samples not covering the timed groups in every round, or a forest whose exact_sums differs from the run's, which the camera-ready suites must set off (pass --exact-sums off for the campaign's acceptance and layout runs). It prints the warnings (sentinel drift, the layout check, G=1 at the survival anchors, E1 against the released numbers).

Re-timing runs on each machine, from that machine's own execution manifest (the native libraries differ), for every run the paper reads: camera-ready, camera-ready-ablation-1, camera-ready-ablation-2 and scenario-v1:

for run in camera-ready camera-ready-ablation-1 camera-ready-ablation-2 scenario-v1; do
  uv run --group baselines treewalker-exp retime $run   # this machine; --manifest <artifacts>/manifests/SUITE.json
done
SYSTEM_INFO=PATH experiments/data/retime/RUN-retime-rN-ARCH.sh   # each script it writes

retime lists every cell and mode where a method's speedup over TreeWalker, or a research variant's ablation effect over its own denominator (the research build's all-on; a monotonic variant's precompute-off counterpart), has a paired interval wider than ±2% of the ratio. It checks the manifest against the run's record (suite, libraries, artifacts, source, settings, cell keys), and writes manifests that time those cells again with twice the rounds, with a script per manifest that runs it as the campaign does (treewalker-exp retime-run: core 0, the PMU required, the system facts recorded, and refused on any VM but the parent run's) into a run of its own, RUN-retime-rN-ARCH. The analysis checks each attempt against its parent and reports each cell's longer attempt; both stay in the data.

A campaign VM that stops is restarted by hand (automatic_restart is off), and startup.sh then picks up where it was. It records every run it launches in ~bench/runs-claimed.txt; on a restart a claimed run resumes by run ID while it is raw, whether or not an invocation of the runner on it finished (a run none finished, partial, is never packed), one already packed is skipped (never timed again, no sentinel appended, not packed twice), and a packed directory that also holds raw cells, or an incomplete pack, stops the boot with nothing deleted (treewalker-exp run-state names the state). An existing directory this deployment never claimed is refused, and run and retime-run refuse a packed run directory themselves. A new boot moves the previous boot's DONE or ERROR into the results' MARKER_HISTORY and out of the bucket, so check_done reads RUNNING until the boot ends. It ends with DONE only if nothing failed: a run with failed cells, a pack or re-time error, or a claimed run that is not complete at the end (treewalker-exp run-complete: packed, every cell of its manifest finished, no failed cell or sentinel attempt unresolved, whichever boot ran it) is recorded, the later runs still go, the results are uploaded, and the boot ends with ERROR naming each failure. Re-time runs are packed only after the re-time phase has resumed them.

On the campaign's VMs re-timing happens on its own: with retime = true (infra/camera-ready.tfvars), startup.sh packs the suites, runs retime for each of the paper's runs, runs every script it writes, packs the re-time runs and uploads them with the rest (retime/ holds the manifests and scripts). By hand, on the VM that ran the campaign (as the bench user, from the repository):

for run in camera-ready camera-ready-ablation-1 camera-ready-ablation-2 scenario-v1; do
  uv run --group baselines treewalker-exp retime $run --out experiments/data/retime
done
for script in experiments/data/retime/*-$(uname -m).sh; do
  SYSTEM_INFO=$HOME/system_info.json sh "$script"   # retime-run: taskset -c 0, --require-pmu
done
for run in experiments/data/runs/*-retime-*/; do
  uv run --group baselines treewalker-exp pack "$run" --remove
done
gcloud storage cp -r experiments/data/runs experiments/data/retime gs://BUCKET/results/MACHINE/

Re-time only on the VM that ran the campaign, and keep it up until the re-times finish. The analysis accepts an attempt only if it ran the parent's program on the parent's protocol and VM: the same executed source and sweep_bench binary (by hash), Rust toolchain and flags, pinned CPUs, PMU, turbo mode, machine type, image, zone and instance (its ID and host name, which the camera-ready runs record), native libraries and tl2cgen runtime, resolved run settings other than the rounds, and per cell the same key, model, input and oracle hashes and compiled baselines. A refused attempt leaves the original reported, marked "re-time refused", and tables lists every attempt and why it was refused (App. A.4).

What to expect from a rerun

Work counters are deterministic: the final run's counts are identical on both machines and to an earlier run's, and on the paper's models they equal the paper's (Table 1's FLCHAIN row, for one). Timings vary with the host: between that earlier run and the final one, on fresh VMs of the same types, TreeWalker's median time per row moved by 1.3% on Intel and 0.3% on Arm. Each cell stops at a 1% precision target, or its block cap or time budget, and reports bootstrap intervals; the sentinel measures the machine's drift over the run, which numbers reports (App. A.6). Exact values move between reruns.

Correctness

A prediction has three stages: tree_sum, the sum of the row's leaf values; raw_margin, tree_sum / divisor + base_score; and the output, after the identity or the sigmoid. When Forest::exact_sums() is true, as it is for the 12 models of the cells measured below, the leaves are added as exact fixed-point integers and rounded once: tree_sum is the f64 nearest to the true sum, and tree order, layout and every ablation except disable_exact_sums leave it unchanged. The margin and the link are ordinary rounded f64 operations.

LightGBM (f64) predictions match Treelite GTIL within 1e-13, and XGBoost (f32) predictions match native XGBoost within 1e-5 (experiments/benchmarks/tests/correctness.rs). GTIL and the full walk add leaves in f64 in tree order, so they differ from the exact sum by their rounding, which grows with the partial sums. Measured on 2026-10-04, the full walk's outputs differ by at most 2.2e-15 on the FLCHAIN and SUPPORT LightGBM cells at horizons 16 and 128, 7.2e-16 on Expedia and 1.1e-14 on FLCHAIN at horizon 1,024; the XGBoost models' outputs are bit-identical.

Citation

@inproceedings{karatay2026treewalker,
  title     = {{TreeWalker}: Partial Evaluation for Grouped Tree-Ensemble Inference},
  author    = {Karatay, Durmus and Newman, Richard},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2026}
}

License

MIT, see LICENSE. This covers the code and the released measurements; the datasets keep their own terms (see Datasets).

About

No description, website, or topics provided.

Resources

Code of conduct

Security policy

Stars

2 stars

Watchers

0 watching

Forks

Releases

Packages

Used by

Contributors

Languages