WIP: FEA Quantile regression for decision trees

doc/modules/tree.rst

@@ -634,6 +634,20 @@ Mean Absolute Error:
 
 Note that it is 3–6× slower to fit than the MSE criterion as of version 1.8.
 
+Quantile (pinball loss):
+
+.. math::
+
+    q_{\alpha}(y)_m = \underset{y \in Q_m}{\mathrm{quantile}_{\alpha}}(y)
+
+    H(Q_m) = \frac{1}{n_m} \sum_{y \in Q_m}
+    \left(\alpha \max(y - q_{\alpha}(y)_m, 0) +
+    (1-\alpha) \max(q_{\alpha}(y)_m - y, 0)\right)
+
+Use ``criterion="quantile"`` together with the ``quantile`` parameter to
+choose :math:`\alpha \in (0, 1)`. The special case ``quantile=0.5`` corresponds
+to the median.
+
 .. _tree_missing_value_support:
 
 Missing Values Support

doc/whats_new/upcoming_changes/sklearn.tree/32903.feature.rst

@@ -0,0 +1,6 @@
+- :class:`tree.DecisionTreeRegressor`, :class:`tree.ExtraTreeRegressor`,
+  :class:`ensemble.RandomForestRegressor`, and :class:`ensemble.ExtraTreesRegressor`
+  now support `criterion="quantile"` together with the `quantile` parameter to
+  optimize the pinball loss (also known as the quantile loss). This effectively
+  allows to do quantile regression.
+  By :user:`Arthur Lacote <cakedev0>`

sklearn/ensemble/_forest.py

@@ -345,7 +345,10 @@ def fit(self, X, y, sample_weight=None):
         # will raise an error if the underlying tree base estimator can't handle missing
         # values. Only the criterion is required to determine if the tree supports
         # missing values.
-        estimator = type(self.estimator)(criterion=self.criterion)
+        estimator_kwargs = {"criterion": self.criterion}
+        if self.criterion == "quantile":
+            estimator_kwargs["quantile"] = self.quantile
+        estimator = type(self.estimator)(**estimator_kwargs)
         missing_values_in_feature_mask = (
             estimator._compute_missing_values_in_feature_mask(
                 X, estimator_name=self.__class__.__name__
@@ -697,7 +700,10 @@ def __sklearn_tags__(self):
         tags = super().__sklearn_tags__()
         # Only the criterion is required to determine if the tree supports
         # missing values
-        estimator = type(self.estimator)(criterion=self.criterion)
+        estimator_kwargs = {"criterion": self.criterion}
+        if self.criterion == "quantile":
+            estimator_kwargs["quantile"] = self.quantile
+        estimator = type(self.estimator)(**estimator_kwargs)
         tags.input_tags.allow_nan = get_tags(estimator).input_tags.allow_nan
         return tags
 
@@ -1609,14 +1615,16 @@ class RandomForestRegressor(ForestRegressor):
            The default value of ``n_estimators`` changed from 10 to 100
            in 0.22.
 
-    criterion : {"squared_error", "absolute_error", "poisson"}, default="squared_error"
+    criterion : {"squared_error", "absolute_error", "quantile", "poisson"}, \
+            default="squared_error"
         The function to measure the quality of a split. Supported criteria
         are "squared_error" for the mean squared error, which is equal to
         variance reduction as feature selection criterion and minimizes the L2
         loss using the mean of each terminal node, "absolute_error" for the mean
         absolute error, which minimizes the L1 loss using the median of each terminal
-        node, and "poisson" which uses reduction in Poisson deviance to find splits,
-        also using the mean of each terminal node.
+        node, "quantile" which minimizes the pinball loss using the quantile of each
+        terminal node (controlled by ``quantile``), and "poisson" which uses reduction
+        in Poisson deviance to find splits, also using the mean of each terminal node.
 
         .. versionadded:: 0.18
            Mean Absolute Error (MAE) criterion.
@@ -1627,6 +1635,9 @@ class RandomForestRegressor(ForestRegressor):
         .. versionchanged:: 1.9
             Criterion `"friedman_mse"` was deprecated.
 
+        .. versionadded:: 1.9
+           Quantile/Pinball loss criterion
+
     max_depth : int, default=None
         The maximum depth of the tree. If None, then nodes are expanded until
         all leaves are pure or until all leaves contain less than
@@ -1786,6 +1797,10 @@ class RandomForestRegressor(ForestRegressor):
 
         .. versionadded:: 1.4
 
+    quantile : float, default=0.5
+        The quantile to predict when ``criterion="quantile"``. It must be strictly
+        between 0 and 1.
+
     Attributes
     ----------
     estimator_ : :class:`~sklearn.tree.DecisionTreeRegressor`
@@ -1913,6 +1928,7 @@ def __init__(
         ccp_alpha=0.0,
         max_samples=None,
         monotonic_cst=None,
+        quantile=0.5,
     ):
         super().__init__(
             estimator=DecisionTreeRegressor(),
@@ -1929,6 +1945,7 @@ def __init__(
                 "random_state",
                 "ccp_alpha",
                 "monotonic_cst",
+                "quantile",
             ),
             bootstrap=bootstrap,
             oob_score=oob_score,
@@ -1959,6 +1976,7 @@ def __init__(
         self.min_impurity_decrease = min_impurity_decrease
         self.ccp_alpha = ccp_alpha
         self.monotonic_cst = monotonic_cst
+        self.quantile = quantile
 
 
 class ExtraTreesClassifier(ForestClassifier):
@@ -2378,21 +2396,26 @@ class ExtraTreesRegressor(ForestRegressor):
            The default value of ``n_estimators`` changed from 10 to 100
            in 0.22.
 
-    criterion : {"squared_error", "absolute_error", "poisson"}, default="squared_error"
+    criterion : {"squared_error", "absolute_error", "quantile", "poisson"}, \
+            default="squared_error"
         The function to measure the quality of a split. Supported criteria
         are "squared_error" for the mean squared error, which is equal to
         variance reduction as feature selection criterion and minimizes the L2
         loss using the mean of each terminal node, "absolute_error" for the mean
         absolute error, which minimizes the L1 loss using the median of each terminal
-        node, and "poisson" which uses reduction in Poisson deviance to find splits,
-        also using the mean of each terminal node.
+        node, "quantile" which minimizes the pinball loss using the quantile of each
+        terminal node (controlled by ``quantile``), and "poisson" which uses reduction
+        in Poisson deviance to find splits, also using the mean of each terminal node.
 
         .. versionadded:: 0.18
            Mean Absolute Error (MAE) criterion.
 
         .. versionchanged:: 1.9
             Criterion `"friedman_mse"` was deprecated.
 
+        .. versionadded:: 1.9
+           Quantile/Pinball loss criterion
+
     max_depth : int, default=None
         The maximum depth of the tree. If None, then nodes are expanded until
         all leaves are pure or until all leaves contain less than
@@ -2556,6 +2579,10 @@ class ExtraTreesRegressor(ForestRegressor):
 
         .. versionadded:: 1.4
 
+    quantile : float, default=0.5
+        The quantile to predict when ``criterion="quantile"``. It must be strictly
+        between 0 and 1.
+
     Attributes
     ----------
     estimator_ : :class:`~sklearn.tree.ExtraTreeRegressor`
@@ -2667,6 +2694,7 @@ def __init__(
         ccp_alpha=0.0,
         max_samples=None,
         monotonic_cst=None,
+        quantile=0.5,
     ):
         super().__init__(
             estimator=ExtraTreeRegressor(),
@@ -2683,6 +2711,7 @@ def __init__(
                 "random_state",
                 "ccp_alpha",
                 "monotonic_cst",
+                "quantile",
             ),
             bootstrap=bootstrap,
             oob_score=oob_score,
@@ -2713,6 +2742,7 @@ def __init__(
         self.min_impurity_decrease = min_impurity_decrease
         self.ccp_alpha = ccp_alpha
         self.monotonic_cst = monotonic_cst
+        self.quantile = quantile
 
 
 class RandomTreesEmbedding(TransformerMixin, BaseForest):

sklearn/ensemble/_gb.py

@@ -378,6 +378,7 @@ class BaseGradientBoosting(BaseEnsemble, metaclass=ABCMeta):
     }
     _parameter_constraints.pop("splitter")
     _parameter_constraints.pop("monotonic_cst")
+    _parameter_constraints.pop("quantile")
 
     @abstractmethod
     def __init__(

sklearn/ensemble/tests/test_forest.py

@@ -294,12 +294,12 @@ def test_probability(name):
     "name, criterion",
     itertools.chain(
         product(FOREST_CLASSIFIERS, ["gini", "log_loss"]),
-        product(FOREST_REGRESSORS, ["squared_error", "absolute_error"]),
+        product(FOREST_REGRESSORS, ["squared_error", "absolute_error", "quantile"]),
     ),
 )
 def test_importances(dtype, name, criterion):
     tolerance = 0.01
-    if name in FOREST_REGRESSORS and criterion == "absolute_error":
+    if name in FOREST_REGRESSORS and criterion in {"absolute_error", "quantile"}:
         tolerance = 0.05
 
     # cast as dtype

sklearn/tree/_classes.py

@@ -76,6 +76,7 @@
     "squared_error": _criterion.MSE,
     "absolute_error": _criterion.MAE,
     "poisson": _criterion.Poisson,
+    "quantile": _criterion.Pinball,
 }
 
 DENSE_SPLITTERS = {"best": _splitter.BestSplitter, "random": _splitter.RandomSplitter}
@@ -382,7 +383,14 @@ def _fit(
                     self.n_outputs_, self.n_classes_
                 )
             else:
-                criterion = CRITERIA_REG[self.criterion](self.n_outputs_, n_samples)
+                args = (self.n_outputs_, n_samples)
+                if self.criterion == "quantile":
+                    args = (*args, self.quantile)
+                if self.criterion == "absolute_error":
+                    # FIXME: this is coupled with code at a much lower level
+                    # because of the inheritance behavior of __cinit__
+                    args = (*args, 0.5)
+                criterion = CRITERIA_REG[self.criterion](*args)
         else:
             # Make a deepcopy in case the criterion has mutable attributes that
             # might be shared and modified concurrently during parallel fitting
@@ -1117,14 +1125,16 @@ class DecisionTreeRegressor(RegressorMixin, BaseDecisionTree):
 
     Parameters
     ----------
-    criterion : {"squared_error", "absolute_error", "poisson"}, default="squared_error"
+    criterion : {"squared_error", "absolute_error", "quantile", "poisson"}, \
+            default="squared_error"
         The function to measure the quality of a split. Supported criteria
         are "squared_error" for the mean squared error, which is equal to
         variance reduction as feature selection criterion and minimizes the L2
         loss using the mean of each terminal node, "absolute_error" for the mean
         absolute error, which minimizes the L1 loss using the median of each terminal
-        node, and "poisson" which uses reduction in Poisson deviance to find splits,
-        also using the mean of each terminal node.
+        node, "quantile" which minimizes the pinball loss using the quantile of each
+        terminal node (controlled by ``quantile``), and "poisson" which uses reduction
+        in Poisson deviance to find splits, also using the mean of each terminal node.
 
         .. versionadded:: 0.18
            Mean Absolute Error (MAE) criterion.
@@ -1135,6 +1145,9 @@ class DecisionTreeRegressor(RegressorMixin, BaseDecisionTree):
         .. versionchanged:: 1.9
             Criterion `"friedman_mse"` was deprecated.
 
+        .. versionadded:: 1.9
+           Quantile/Pinball loss criterion
+
     splitter : {"best", "random"}, default="best"
         The strategy used to choose the split at each node. Supported
         strategies are "best" to choose the best split and "random" to choose
@@ -1255,6 +1268,10 @@ class DecisionTreeRegressor(RegressorMixin, BaseDecisionTree):
 
         .. versionadded:: 1.4
 
+    quantile : float, default=0.5
+        The quantile to predict when ``criterion="quantile"``. It must be strictly
+        between 0 and 1. If 0.5 (default), the model predicts the median.
+
     Attributes
     ----------
     feature_importances_ : ndarray of shape (n_features,)
@@ -1338,9 +1355,10 @@ class DecisionTreeRegressor(RegressorMixin, BaseDecisionTree):
     _parameter_constraints: dict = {
         **BaseDecisionTree._parameter_constraints,
         "criterion": [
-            StrOptions({"squared_error", "absolute_error", "poisson"}),
+            StrOptions({"squared_error", "absolute_error", "poisson", "quantile"}),
             Hidden(Criterion),
         ],
+        "quantile": [Interval(RealNotInt, 0.0, 1.0, closed="neither")],
     }
 
     def __init__(
@@ -1358,6 +1376,7 @@ def __init__(
         min_impurity_decrease=0.0,
         ccp_alpha=0.0,
         monotonic_cst=None,
+        quantile=0.5,
     ):
         if isinstance(criterion, str) and criterion == "friedman_mse":
             # TODO(1.11): remove support of "friedman_mse" criterion.
@@ -1383,6 +1402,7 @@ def __init__(
             ccp_alpha=ccp_alpha,
             monotonic_cst=monotonic_cst,
         )
+        self.quantile = quantile
 
     @_fit_context(prefer_skip_nested_validation=True)
     def fit(self, X, y, sample_weight=None, check_input=True):
@@ -1767,14 +1787,16 @@ class ExtraTreeRegressor(DecisionTreeRegressor):
 
     Parameters
     ----------
-    criterion : {"squared_error", "absolute_error", "poisson"}, default="squared_error"
+    criterion : {"squared_error", "absolute_error", "quantile", "poisson"}, \
+            default="squared_error"
         The function to measure the quality of a split. Supported criteria
         are "squared_error" for the mean squared error, which is equal to
         variance reduction as feature selection criterion and minimizes the L2
         loss using the mean of each terminal node, "absolute_error" for the mean
         absolute error, which minimizes the L1 loss using the median of each terminal
-        node, and "poisson" which uses reduction in Poisson deviance to find splits,
-        also using the mean of each terminal node.
+        node, "quantile" which minimizes the pinball loss using the quantile of each
+        terminal node (controlled by ``quantile``), and "poisson" which uses reduction
+        in Poisson deviance to find splits, also using the mean of each terminal node.
 
         .. versionadded:: 0.18
            Mean Absolute Error (MAE) criterion.
@@ -1785,6 +1807,9 @@ class ExtraTreeRegressor(DecisionTreeRegressor):
         .. versionchanged:: 1.9
             Criterion `"friedman_mse"` was deprecated.
 
+        .. versionadded:: 1.9
+           Quantile/Pinball loss criterion
+
     splitter : {"random", "best"}, default="random"
         The strategy used to choose the split at each node. Supported
         strategies are "best" to choose the best split and "random" to choose
@@ -1897,6 +1922,10 @@ class ExtraTreeRegressor(DecisionTreeRegressor):
 
         .. versionadded:: 1.4
 
+    quantile : float, default=0.5
+        The quantile to predict when ``criterion="quantile"``. It must be strictly
+        between 0 and 1. If 0.5 (default), the model predicts the median.
+
     Attributes
     ----------
     max_features_ : int
@@ -1981,6 +2010,7 @@ def __init__(
         max_leaf_nodes=None,
         ccp_alpha=0.0,
         monotonic_cst=None,
+        quantile=0.5,
     ):
         super().__init__(
             criterion=criterion,
@@ -1995,6 +2025,7 @@ def __init__(
             random_state=random_state,
             ccp_alpha=ccp_alpha,
             monotonic_cst=monotonic_cst,
+            quantile=quantile,
         )
 
     def __sklearn_tags__(self):