From 154e9806f6dce527143a4d48c91400f35bbc47c1 Mon Sep 17 00:00:00 2001 From: Moussa Ouallaf Date: Tue, 3 Nov 2026 12:31:05 +0000 Subject: [PATCH] Check the input once in pred_param and skip trees that never split pred_param called every base tree with its own input check, which costs far more than the tree walk itself. When every learner is a scikit-learn regression tree, pred_param now checks X once (float32, the columns the trees read, NaN allowed only where every tree allows it) and calls the trees with check_input=False. A tree whose leaves all hold one value is replaced by its scalar update. The floating-point operations per parameter happen in the same order as before, so the output is bit-identical. The precomputed plan keeps references to the rounds it was built from and is rebuilt when they, init_params or learning_rate change. It is not pickled. max_iter, sparse input and other learners use the loop as before. fit and partial_fit warn (NGBoostDegenerateWarning) when every learner is constant, which happens silently when int(minibatch_frac * n) is below max(min_samples_split, 2 * min_samples_leaf). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01N8iMyoH5TsKptT8vKngJEo --- ngboost/ngboost.py | 143 +++++++++++++++++++++- tests/test_prediction_plan.py | 224 ++++++++++++++++++++++++++++++++++ 2 files changed, 366 insertions(+), 1 deletion(-) create mode 100644 tests/test_prediction_plan.py diff --git a/ngboost/ngboost.py b/ngboost/ngboost.py index 3d3e6702..69771640 100644 --- a/ngboost/ngboost.py +++ b/ngboost/ngboost.py @@ -4,12 +4,14 @@ # pylint: disable=unused-argument,too-many-locals,too-many-branches,too-many-statements # pylint: disable=unused-variable,invalid-unary-operand-type,attribute-defined-outside-init # pylint: disable=redundant-keyword-arg,protected-access,unnecessary-lambda-assignment +import warnings + import numpy as np from joblib import Parallel, delayed from scipy.sparse import issparse from sklearn.base import clone from sklearn.model_selection import train_test_split -from sklearn.tree import DecisionTreeRegressor +from sklearn.tree import DecisionTreeRegressor, ExtraTreeRegressor from sklearn.utils import check_array, check_random_state, check_X_y from ngboost.distns import MultivariateNormal, Normal, k_categorical @@ -40,6 +42,120 @@ def _validate_sample_weight(sample_weight, X, input_name="sample_weight"): return sample_weight +class NGBoostDegenerateWarning(UserWarning): + """Every base learner is constant, so the model predicts one distribution for all inputs.""" + + +def _constant_tree_value(tree): + """The output of a fitted regression tree whose leaves all hold the same value, else None.""" + leaves = tree.tree_.value[tree.tree_.children_left == -1].reshape(-1) + bits = np.ascontiguousarray(leaves, dtype=np.float64).view(np.uint64) + if tree.n_outputs_ != 1 or not (bits == bits[0]).all(): + return None + return tree.predict(np.zeros((1, tree.n_features_in_)))[0] + + +class _PredictionPlan: + """ + pred_param for a model whose base learners are all scikit-learn regression trees. + + The input is checked once for all trees instead of once per tree, and a tree whose leaves all + hold the same value (it never split) is replaced by its update, a scalar. For each parameter, + the constant updates before its first tree that splits are subtracted from the initial + parameters once; later ones stay as scalar steps in round order. Every parameter therefore goes + through the same floating-point operations in the same order as the boosting loop, and the + result is bit-identical to it. + + The plan keeps references to the rounds it was built from; matches() rejects it once the + model's rounds, initial parameters or learning rate change, so it can never go stale. + """ + + def __init__(self, model, n_params): + self.learning_rate = model.learning_rate + self.init_params = np.array(model.init_params, dtype=float) + self.base_models = [list(models) for models in model.base_models] + self.scalings = list(model.scalings) + self.col_idxs = list(model.col_idxs) + self.start = np.ones(n_params) * model.init_params + self.steps, self.n_constant, self.n_learners = [], 0, 0 + split_seen = np.zeros(n_params, dtype=bool) + for i, (models, s) in enumerate(zip(self.base_models, self.scalings)): + for k, tree in enumerate(models): + self.n_learners += 1 + value = _constant_tree_value(tree) + if value is None: + split_seen[k] = True + self.steps.append((i, k, None)) + continue + self.n_constant += 1 + # the update the boosting loop makes for this parameter + step = self.learning_rate * value * s + if split_seen[k]: + self.steps.append((i, k, step)) + else: + self.start[k] -= step + cols = np.unique( + np.concatenate([np.asarray(c, dtype=np.intp) for c in self.col_idxs]) + ) + self.cols = cols + self.all_cols = np.array_equal(cols, np.arange(cols.size)) + # positions of each round's columns in X[:, cols]; None when a round reads all of them + self.positions = [ + None if np.array_equal(c, cols) else np.searchsorted(cols, c) + for c in self.col_idxs + ] + trees = [tree for models in self.base_models for tree in models] + # the trees' own check allows NaN only where every tree supports missing values + self.allow_nan = all( + tree.__sklearn_tags__().input_tags.allow_nan and tree.monotonic_cst is None + for tree in trees + ) + + @staticmethod + def applies_to(model): + return bool(model.base_models) and all( + type(learner) in (DecisionTreeRegressor, ExtraTreeRegressor) + for models in model.base_models + for learner in models + ) + + def matches(self, model): + return ( + model.learning_rate == self.learning_rate + and np.array_equal(model.init_params, self.init_params) + and len(model.base_models) == len(self.base_models) + and all( + len(new) == len(old) and all(a is b for a, b in zip(new, old)) + for new, old in zip(model.base_models, self.base_models) + ) + and len(model.col_idxs) == len(self.col_idxs) + and all(a is b for a, b in zip(model.col_idxs, self.col_idxs)) + and list(model.scalings) == self.scalings + ) + + def pred_param(self, X): + if not (self.all_cols and X.shape[1] == self.cols.size): + # only the columns the trees read, as the loop's X[:, col_idx] + X = X[:, self.cols] + X = check_array( + X, + dtype=np.float32, + ensure_all_finite="allow-nan" if self.allow_nan else True, + ) + params = np.ones((X.shape[0], self.start.size)) * self.start + last, X_i = None, None + for i, k, step in self.steps: + if step is not None: + params[:, k] -= step + continue + if i != last: + pos = self.positions[i] + last, X_i = i, X if pos is None else X[:, pos] + pred = self.base_models[i][k].predict(X_i, check_input=False) + params[:, k] -= self.learning_rate * pred * self.scalings[i] + return params + + class NGBoost: """ Constructor for all NGBoost models. @@ -146,6 +262,7 @@ def __getstate__(self): state = self.__dict__.copy() # Remove the unpicklable entries. del state["Manifold"] + state.pop("_pred_plan", None) # rebuilt on the first prediction state["Dist_name"] = self.Dist.__name__ if self.Dist.__name__ == "Categorical": del state["Dist"] @@ -175,7 +292,20 @@ def fit_init_params_to_marginal(self, Y, sample_weight=None, iters=1000): Y ) # would be best to put sample weights here too + def _prediction_plan(self): + plan = getattr(self, "_pred_plan", None) + if plan is None or not plan.matches(self): + plan = None + if _PredictionPlan.applies_to(self): + plan = _PredictionPlan(self, self.Manifold.n_params) + self._pred_plan = plan + return plan + def pred_param(self, X, max_iter=None): + if not max_iter and not issparse(X): + plan = self._prediction_plan() + if plan is not None: + return plan.pred_param(X) m, n = X.shape params = np.ones((m, self.Manifold.n_params)) * self.init_params for i, (models, s, col_idx) in enumerate( @@ -598,6 +728,17 @@ def partial_fit( if X_val is not None and Y_val is not None: self.evals_result["val"] = {metric: val_loss_list} + plan = self._prediction_plan() + if plan is not None and plan.n_constant == plan.n_learners: + warnings.warn( + f"All {plan.n_learners} base learners are constant, so the model predicts " + "the same distribution for every input. Each round's trees are fitted on " + "int(minibatch_frac * n) rows and cannot split below " + "max(min_samples_split, 2 * min_samples_leaf) rows.", + NGBoostDegenerateWarning, + stacklevel=2, + ) + return self def _set_base_params(self, parameters): diff --git a/tests/test_prediction_plan.py b/tests/test_prediction_plan.py new file mode 100644 index 00000000..3680d698 --- /dev/null +++ b/tests/test_prediction_plan.py @@ -0,0 +1,224 @@ +"""pred_param checks the input once and skips constant trees, bit-identical to the loop.""" + +# pylint: disable=protected-access +import pickle +import warnings + +import numpy as np +import pytest +from scipy import sparse +from sklearn.linear_model import Ridge +from sklearn.tree import DecisionTreeRegressor, ExtraTreeRegressor + +from ngboost import NGBClassifier, NGBoost, NGBRegressor +from ngboost.distns import Bernoulli, Normal, T, k_categorical +from ngboost.ngboost import NGBoostDegenerateWarning +from ngboost.scores import LogScore + + +def loop_pred_param(model, X, max_iter=None): + """The boosting loop pred_param runs without a plan: the reference for every test.""" + params = np.ones((X.shape[0], model.Manifold.n_params)) * model.init_params + for i, (models, s, col_idx) in enumerate( + zip(model.base_models, model.scalings, model.col_idxs) + ): + if max_iter and i == max_iter: + break + resids = np.array([m.predict(X[:, col_idx]) for m in models]).T + params -= model.learning_rate * resids * s + return params + + +def assert_bitwise_equal(a, b): + a, b = np.asarray(a, dtype=np.float64), np.asarray(b, dtype=np.float64) + np.testing.assert_array_equal(a.view(np.uint64), b.view(np.uint64)) + + +def toy_data(n=300, seed=0): + rng = np.random.default_rng(seed) + X = rng.normal(size=(n, 4)) + Y = 1.5 * (X[:, 0] > 0.3) + rng.normal(scale=0.5 + 0.8 * (X[:, 1] > 0), size=n) + return X, Y + + +class NormalFixedInit(Normal): + """A fixed start, so that constant trees take nonzero steps.""" + + @staticmethod + def fit(Y): + return np.array([0.0, 0.0]) + + +def constant_tree(): + return DecisionTreeRegressor(max_depth=2, min_samples_leaf=10_000, random_state=0) + + +def split_tree(): + return DecisionTreeRegressor(max_depth=2, min_samples_leaf=10, random_state=0) + + +def make_ngb(base, n_estimators=40, **kwargs): + return NGBoost( + Dist=NormalFixedInit, + Score=LogScore, + Base=base, + n_estimators=n_estimators, + learning_rate=0.05, + verbose=False, + random_state=0, + **kwargs, + ) + + +def fit_quietly(model, X, Y): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", NGBoostDegenerateWarning) + return model.fit(X, Y) + + +def test_all_constant_model_warns_and_calls_no_tree(): + X, Y = toy_data() + with pytest.warns(NGBoostDegenerateWarning): + model = make_ngb(constant_tree(), minibatch_frac=0.7, col_sample=0.5).fit(X, Y) + assert np.abs(model.pred_param(X)[0] - model.init_params).max() > 0.1 + assert_bitwise_equal(model.pred_param(X), loop_pred_param(model, X)) + assert not model._pred_plan.steps + + +def test_no_warning_when_trees_split(): + X, Y = toy_data() + with warnings.catch_warnings(): + warnings.simplefilter("error", NGBoostDegenerateWarning) + model = make_ngb(split_tree(), n_estimators=5).fit(X, Y) + assert_bitwise_equal(model.pred_param(X), loop_pred_param(model, X)) + + +def test_constant_and_split_trees_per_parameter(): + X, Y = toy_data() + model = make_ngb([split_tree(), constant_tree()]).fit(X, Y) + assert_bitwise_equal(model.pred_param(X), loop_pred_param(model, X)) + assert {k for _, k, step in model._pred_plan.steps if step is None} == {0} + + +def test_plan_follows_edited_rounds(): + X, Y = toy_data() + model = make_ngb([split_tree(), split_tree()], n_estimators=30).fit(X, Y) + model.pred_param(X) + for models in model.base_models[10:]: # constant trees after trees that split + for k in range(2): + models[k] = constant_tree().fit(X, np.full(len(X), 0.01 * (k + 1))) + assert_bitwise_equal(model.pred_param(X), loop_pred_param(model, X)) + assert model._pred_plan.n_constant == 40 + model.base_models = model.base_models[:7] + model.scalings = model.scalings[:7] + model.col_idxs = model.col_idxs[:7] + assert_bitwise_equal(model.pred_param(X), loop_pred_param(model, X)) + model.learning_rate = 0.1 + assert_bitwise_equal(model.pred_param(X), loop_pred_param(model, X)) + + +def test_max_iter_and_staged_predictions(): + X, Y = toy_data() + model = make_ngb([split_tree(), constant_tree()], n_estimators=20).fit(X, Y) + for max_iter in (1, 5, 19): + assert_bitwise_equal( + model.pred_param(X, max_iter=max_iter), + loop_pred_param(model, X, max_iter=max_iter), + ) + assert_bitwise_equal( + model.staged_pred_dist(X)[-1]._params, model.pred_dist(X)._params + ) + + +def test_refit_and_partial_fit(): + X, Y = toy_data() + X2, Y2 = toy_data(seed=1) + model = fit_quietly(make_ngb(constant_tree(), n_estimators=15), X, Y) + fit_quietly(model, X2, Y2 + 3.0) + assert_bitwise_equal(model.pred_param(X2), loop_pred_param(model, X2)) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", NGBoostDegenerateWarning) + model.partial_fit(X, Y) + assert len(model.base_models) == 30 + assert_bitwise_equal(model.pred_param(X), loop_pred_param(model, X)) + + +def test_plan_is_not_pickled(): + X, Y = toy_data() + model = fit_quietly(make_ngb(constant_tree(), n_estimators=15), X, Y) + loaded = pickle.loads(pickle.dumps(model)) + assert getattr(loaded, "_pred_plan", None) is None + assert_bitwise_equal(loaded.pred_param(X), loop_pred_param(model, X)) + + +@pytest.mark.parametrize("dist", [Normal, T]) +def test_regressor_with_early_stopping(dist): + X, Y = toy_data() + model = NGBRegressor( + Dist=dist, + Base=DecisionTreeRegressor(max_depth=2, min_samples_leaf=60), + n_estimators=30, + verbose=False, + random_state=0, + minibatch_frac=0.5, + col_sample=0.5, + early_stopping_rounds=3, + ).fit(X, Y) + assert_bitwise_equal(model.pred_param(X), loop_pred_param(model, X)) + + +@pytest.mark.parametrize("dist,K", [(Bernoulli, 2), (k_categorical(3), 3)]) +def test_classifiers(dist, K): + X, Y = toy_data() + y = np.digitize(Y, np.quantile(Y, np.linspace(0, 1, K + 1)[1:-1])) + model = NGBClassifier( + Dist=dist, Base=split_tree(), n_estimators=20, verbose=False, random_state=0 + ).fit(X, y) + assert_bitwise_equal(model.pred_param(X), loop_pred_param(model, X)) + + +@pytest.mark.parametrize( + "base", + [ + ExtraTreeRegressor(max_depth=3, random_state=0), + DecisionTreeRegressor(max_depth=3, monotonic_cst=[1, 0, 0, 0]), + ], +) +def test_other_tree_settings(base): + X, Y = toy_data() + model = NGBRegressor(Base=base, n_estimators=20, verbose=False, random_state=0) + model.fit(X, Y) + assert_bitwise_equal(model.pred_param(X), loop_pred_param(model, X)) + + +def test_missing_values_follow_the_trees(): + X, Y = toy_data() + X[::7, 2] = np.nan + model = NGBRegressor(Base=split_tree(), n_estimators=20, verbose=False).fit(X, Y) + assert_bitwise_equal(model.pred_param(X), loop_pred_param(model, X)) + no_nan = DecisionTreeRegressor(max_depth=2, criterion="absolute_error") + model = NGBRegressor(Base=no_nan, n_estimators=5, verbose=False) + model.fit(*toy_data()) + with pytest.raises(ValueError): + model.pred_param(X) + + +def test_input_width(): + X, Y = toy_data() + model = fit_quietly(make_ngb(constant_tree(), n_estimators=10), X, Y) + wide = np.hstack([X, X]) + assert_bitwise_equal(model.pred_param(wide), loop_pred_param(model, wide)) + with pytest.raises(IndexError): + model.pred_param(X[:, :2]) + with pytest.raises(ValueError): + model.pred_param(np.where(np.arange(4) == 1, np.inf, X)) + + +def test_other_learners_and_sparse_input_use_the_loop(): + X, Y = toy_data() + model = NGBRegressor(Base=Ridge(), n_estimators=10, verbose=False).fit(X, Y) + assert_bitwise_equal(model.pred_param(X), loop_pred_param(model, X)) + assert model._pred_plan is None + model = NGBRegressor(Base=split_tree(), n_estimators=10, verbose=False).fit(X, Y) + X_sparse = sparse.csr_matrix(X) + assert_bitwise_equal(model.pred_param(X_sparse), loop_pred_param(model, X_sparse))