Tags: NVlabs/gbrl
Tags
Fix optimizer-aware SHAP completeness; add local Adam TreeSHAP (#27) * Fix SHAP local accuracy for SGD; warn on Adam (v1.1.9) SHAP was decomposing raw gradient leaf values instead of the optimizer-adjusted contributions that predict() returns, violating local accuracy: E[predict] + sum_f(shap_f(x)) != predict(x). SGD fix: scale each tree's leaf predictions by -lr_t inside tree_shap() and ensemble_shap() before running LinearTreeSHAP. This makes completeness hold to floating-point precision for both constant and linear lr schedules, and for multi-output models. Adam: issue a RuntimeWarning when shap()/tree_shap() is called on an Adam-optimized model. Adam applies an adaptive per-sample step whose contribution is not a function of the leaf value alone, so the current raw-gradient decomposition does not reconstruct the model prediction. Full optimizer-aware Adam SHAP is planned. Tests: replace the sklearn-comparison assertion in test_shap_cpu with a proper completeness check (E[pred] + sum(tree_shap) == pred). Add test_ensemble_shap_completeness_sgd covering multi-tree, multi-output, and multiple learning rates. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> * Add optimizer-aware SHAP for SGD and Adam (v1.1.9) SHAP was decomposing raw gradient leaf values instead of the actual optimizer contributions that predict() returns, violating local accuracy: E[predict] + sum_f(shap_f(x)) != predict(x). SGD fix: scale each tree's leaf predictions by -lr_t before running LinearTreeSHAP. Completeness holds to floating-point precision for constant and linear lr schedules, and for multi-output models. Adam fix: for each sample and tree, freeze the pre-tree Adam moment state (m, v), compute the one-step Adam delta for every leaf in that tree under that state, and use those per-sample effective values in TreeSHAP. This gives exact local accuracy per sample per tree ("one-step surrogate" completeness). tree_shap() replays trees 0..t-1 first to recover the correct pre-tree moment state. Implementation: - shap.h/cpp: add node_to_leaf_idx and leaf_cond_probs to shapData so leaf identities and edge-weight products are available without re-traversing the tree - gbrl.cpp: three new static helpers (apply_optimizer_shap_predictions, find_factual_leaf, advance_adam_state) + refactored tree_shap / ensemble_shap to branch on has_adam; SGD path unchanged in cost, Adam path runs sample-inner loop with per-sample moment replay - Removed interim Adam RuntimeWarning from gbt_learner.py Tests: completeness checks for SGD (lr=0.1/0.05, output_dim=1/2) and Adam (lr=0.01, output_dim=1/2) added to test_gbt_single.py. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> * Fix overlapping SGD ranges, fix base-value comments, add linear-lr test Address PR review findings: 1. Overlapping SGD optimizer ranges: the old per-optimizer assignment would silently drop earlier contributions when two SGD optimizers cover the same output dimension. Pre-compute a per-dimension scale (summing -lr over all covering SGD optimizers) and assign once, matching what predict() does (sequential step() calls are additive). 2. Comment/docstring accuracy: replace "bias + sum(shap) == predict" with "base + sum(shap) == predict" throughout, making clear that the base value is E_x[predict(x)] (the explainer's expected value), not the model bias parameter. 3. Linear lr scheduler coverage: extend test_ensemble_shap_completeness_sgd to train with a LinearScheduler (init_lr=0.1 → stop_lr=0.01 over 30 steps) and assert completeness, catching any off-by-one in the tree_idx passed to scheduler->get_lr(). Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> * Add base-value API, reject overlapping optimizer ranges, strengthen Adam test Address four PR review blockers: 1. Overlapping optimizer output ranges: set_optimizer() now throws std::runtime_error if a new optimizer's [start_idx, stop_idx) overlaps any existing range, matching the additive semantics of predict() and preventing silently wrong SHAP scaling for mixed/overlapping configs. 2. Base value API: ensemble_shap() and tree_shap() accept an optional float *base_values parameter; when non-null they accumulate the sample-specific SHAP base (bias + sum_t sum_l cond_prob_l * delta_l). New bindings ensemble_shap_and_base / tree_shap_and_base return (phi, base) as a Python tuple. Python shap()/tree_shap() gain return_base=False; when True they call the new bindings so callers can verify base[s] + phi[s].sum(axis=0) == predict(x_s) exactly. For Adam the base is sample-specific; for SGD it is shared. 3. Stronger Adam test: test_ensemble_shap_completeness_adam now verifies (a) full ensemble completeness with the returned sample-specific base to 1e-4, exercising the replayed moment-state logic across all trees; (b) that tree_shap(1) base varies across samples (proving Adam state from tree 0 influences tree 1's effective leaf values); (c) that per-tree contributions are non-trivial. 4. Docstrings: shap() and tree_shap() in gbt_learner.py now document Adam as local optimizer-aware TreeSHAP (not exact global SHAP), explain what the base value represents, and note that tree_shap(t) for Adam explains tree t under the factual moment state from all preceding trees. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> * fixing some bugs * fixing unittests * change test tolerance * fixing pava issue * fixing pava * fixing bugs * fixing docstring * CR fixes * fixing CR * fixed a unit test * fixing CR * fixing runtime warning * fixing issues * fixing CR * fixed isotonic regression * fixed more stuff * fixing more CR * fixing mistakes * more bugs * fix issue * latest fixes * fixing more bugs * fixing more bugs in monototnic constraints * more fixes * fixing latest issues * fixing more bugs * adding more checks * latest fixes * fixed latest * fixes * fixing more bugs * fixred more issues --------- Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
PreviousNext