Optimal Sparse Recovery with Decision Stumps

Decision trees are one of the most popular tools used in machine learning. Due to their simplicity and interpretability, trees are widely implemented by data scientist, both individually, and in aggregation with ensemble methods such as random forests and gradient boosting (Friedman 2001).

In addition to their predictive accuracy, tree based methods are an important tool used for the variable selection problem: identifying a relevant small subset of a high-dimensional feature space of the input variables that can accurately predict the output. When the relationship between the variables is linear, it has long been known that Lasso achieves the optimal sample complexity rate for this problem (Wainwright 2009a). In practice, however, tree-based methods have been shown to be preferable as they scale linearly with the size of the data and can capture non-linear relationships between the variables (Xu et al. 2014).

Notably, various tree structured systems are implemented for this variable selection task across wide-spanning domains. For example, tree based importance measures have been used for prediction of financial distress (Qian et al. 2022), wireless signal recognition (Li and Wang 2018), credit scoring (Xia et al. 2017), and the discovery of genes in the field of bioinformatics (Breiman 2001; Bureau et al. 2005; Huynh-Thu et al. 2012; Lunetta et al. 2004) to name a small fraction.

Despite this empirical success however, the theoretical properties of these tree based methods for feature selection are not well understood. While several papers have considered variants of this problem (see Section 3 for an overview), even in the simple linear case, the sample complexity of the decision tree is not well-characterized.

In this paper, we attempt to bridge this gap and analyze the variable selection properties of single-level decision trees, commonly referred to as “decision stumps” (DStump ). Considering both linear and non-linear settings, we show that DStump achieves nearly-tight sample complexity rates for a variety of practical sample distributions. Compared to prior work, our analysis is simpler and applies to different variants of the decision tree, as well as more general function classes.

The remainder of the paper is organized as follows: in Section 2 we introduce our main results on the finite sample guarantees of DStump , in Section 3 we discuss important prior results in the field of sparse recovery and where they fall flat, in Section 4 we present the recovery algorithm, and in Section 5, we provide theoretical guarantees for the procedure under progressively more general problem assumptions. The proofs of these results provided in Section 6. We supplement the theoretical results with computational experiments in Section 7 and finally conclude the paper in Section 8.

We assume that we are given a dataset $\mathcal{D}=\{(\mathbf{X}_{i,:},Y_{i})\}_{i=1}^{n}$ consisting of $n$ samples from the non-parametric regression model $Y_{i}=\sum_{k}f_{k}(\mathbf{X}_{i,k})+W_{i}$ where $i$ denotes the sample number, $\mathbf{X}_{i,:}\in\mathbb{R}^{p}$ is the input vector with corresponding output $Y_{i}\in\mathbb{R}$, $W_{i}\in\mathbb{R}$ are i.i.d noise values and $f_{k}:\mathbb{R}\to\mathbb{R}$ are univariate functions that are $s$-sparse:

For the recovery problem, we consider the DStump algorithm as first proposed by (Kazemitabar et al. 2017). The algorithm, shown in Algorithm 1, fits a single-level decision tree (stump) on each feature using either the “median split” or the “optimal split” strategy and ranks the features by the error of the corresponding trees. The median split provides a substantially simplified implementation as we do not need to optimize the stump construction, further providing an improved run time over the comparable optimal splitting procedure. Indeed, the median and optimal split have time complexity at most $\mathcal{O}(np)$ and $\mathcal{O}(np\log(n))$ respectively.

In spite of this simplification, we show that in the widely considered case of linear design, where the $f_{k}$ are linear, DStump can correctly recover $S$ with a sample complexity bound of $\mathcal{O}(s\log p)$, matching the minimax optimal lower bound for the problem as achieved by Lasso (Wainwright 2009a, b). This result is noteworthy and surprising since the DStump algorithm (and decision trees in general) is not designed with a linearity assumption, as is the case with Lasso . For this reason, trees are in general utilized for their predictive power in a non-linear model, yet our work proves their value in the opposite. We further extend these results for non-linear models and general sub-Gaussian distributions, improving previously known results using simpler analysis. In addition, our results do not require the sparsity level $s$ to be known in advance. We summarize our main technical results as follows:

• •

  We obtain a sample complexity bound of $\mathcal{O}(s\log p)$ for the DStump algorithm in the linear design case, matching the optimal rate of Lasso and improving prior bounds in the literature for both the median and optimal split. This is the first tight bound on the sample complexity of single-depth decision trees used for sparse recovery and significantly improves on the existing results.

• •

  We extend our results to the case of non-linear $f_{k}$, obtaining tighter results for a wider class of functions compared to the existing literature. We further use these results to generalize our analysis to sub-Gaussian distributions via the extension to nonlinear $f_{k}$.

• •

  As a byproduct of our improved analysis, we show that our results hold for the case where the number of active features $s$ is not known. This is the first theoretical guarantee on decision stumps that does not require $s$ to be known.

• •

  We validate our theoretical results using numerical simulations that show the necessity of our analytic techniques.

While our model framework and the sparsity problem as a whole have been studied extensively (Fan and Lv 2006; Lafferty and Wasserman 2008; Wainwright 2009a, b), none have replicated the well known optimal lower bound for the classification problem under the given set of assumptions. Our work provides improved finite sample guarantees on DStump for the regression problem that nearly match that of Lasso by means of weak learners.

Most closely related to our work is that of (Kazemitabar et al. 2017), which formulates the DStump algorithm and theoretical approach for finite sample guarantees of these weak learners. Unlike our nearly tight result on the number of samples required for recovery, (Kazemitabar et al. 2017) provides a weaker $\mathcal{O}(s^{2}\log p)$ bound when using the median splitting criterion. We here demonstrate that the procedure can obtain the near optimal finite sample guarantees by highlighting a subtle nuance in the analysis of the stump splitting (potentially of independent interest to the reader); instead of using the variance of one sub-tree as an impurity measure, we use the variance of both sub-trees. As we will show both theoretically and experimentally, this consideration is vital for obtaining tight bounds. Our analysis is also more general than that of (Kazemitabar et al. 2017), with applications to both median and optimal splits, a wider class of functions $f_{k}$, and more general distributions.

In a recent work, (Klusowski and Tian 2021) provide an indirect analysis of the DStump formulation with the optimal splitting criterion, by studying its relation to the SIS algorithm of (Fan and Lv 2006). Designed for linear models specifically, SIS sorts the features based on their Pearson correlation with the output, and has optimal sample complexity for the linear setting. (Klusowski and Tian 2021) show that when using the optimal splitting criterion, DStump is equivalent to SIS up to logarithmic factors, which leads to a sample complexity of

This improves the results of (Kazemitabar et al. 2017), though the analysis is more involved. A similar technique is also used to study the non-linear case, but the conditions assumed on $f_{k}$ are hard to interpret and the bounds are weakened. Indeed, instantiating the non-parametric results for the linear case implies a sub-optimal sample complexity rate of $\mathcal{O}(s^{2}\log p)$. In addition, (Klusowski and Tian 2021) describe a heuristic algorithm for the case of unknown $|S|$, though they fail to prove any guarantees on its performance.

In contrast, we provide a direct analysis of DStump . This allows us to obtain optimal bounds for both the median and optimal split in the linear case, as well as improved and generalized bounds in the non-linear case. Our novel approach further allows us to analyze the case of unknown sparsity level, as we provide the first formal proof for the heuristic algorithm suggested in (Klusowski and Tian 2021) and (Fan, Feng, and Song 2011). Compared to prior work, our analysis is considerably simpler and applies to more general settings.

Additionally, various studies have leveraged the simplicity of median splits in decision trees to produce sharp bounds on mean-squared error for the regression problem with random forests (Duroux, Roxane and Scornet, Erwan 2018; Klusowski 2021). In each of these studies, analysis under the median split assumption allows for improvements in both asymptotic and finite sample bounds on prediction error for these ensemble weak learners. In the present work, we extend this intuition to the feature selection problem for high-dimensional sparse systems, and further emphasize the utility of the median split even in the singular decision stump case.

In this section, we prove Theorem 5.3 as it is the more general version of Theorem 5.1, and defer with remainder of the proofs to the appendix.

To prove that the algorithm succeeds, we need to show that $\textnormal{imp}_{k}<\textnormal{imp}_{k^{\prime}}$ for all $k\in S,k^{\prime}\notin S$. We proceed by first proving an upper bound $\textnormal{imp}_{k}$ for all $k\in S$

We now focus on proving Lemma 6.1. We assume without loss of generality that $f_{k}$ is increasing ¹¹1The case of decreasing $f_{k}$ follows by a symmetric argument or by mapping $X^{k}\to-X^{k}$. We further assume that $\mathbb{E}\left[f_{k}(X^{k}_{i})\right]=0$ as DStump is invariant under constant shifts of the output. Finally, we assume that $n>5$, as for $n\leq 5$, the theorem’s statement can be made vacuous by choosing large $c$.  
We will assume $\{n_{L},n_{R}\}=\{\lfloor\frac{n}{2}\rfloor,\lceil\frac{n}{2}\rceil\}$ throughout the proof; as such our results will hold for both the optimal and median splitting criteria. As noted before, a key point for obtaining a tight bound is considering both sub-trees in the analysis instead of considering them individually. Formally, while the impurity is usually defined via variance as in (1), it has the following equivalent definition.

The above identity is commonly used for the analysis of decision trees and their properties (Breiman et al. 1983; Li et al. 2019; Klusowski 2020; Klusowski and Tian 2021). From an analytic perspective, the key difference between (3) and (1) is that the empirical averaging is calculated before taking the square, allowing us to more simply analyze the first moment rather than the second.

Intuitively, we want to use concentration inequalities to show that $\widehat{\mathbb{E}}\left[Y^{k}_{L}\right]$ and $\widehat{\mathbb{E}}\left[Y^{k}_{R}\right]$ concentrate around their expectations and lower bound $|\mathbb{E}\left[Y^{k}_{L}\right]-\mathbb{E}\left[Y^{k}_{R}\right]|$. This is challenging however as concentration results typically require an i.i.d assumption but, as we will see, $Y^{k}_{L,i}$ are not i.i.d. More formally, for each $k$, define the random variable $X^{k}_{L}\in\mathbb{R}^{n_{L}}$ as $(X^{k}_{\tau^{k}(1)},\dots,X^{k}_{\tau^{k}(n_{L})})^{T}$ and thus $Y^{k}_{L,i}=f_{k}(X^{k}_{L,i})+\sum_{j\neq k}f_{j}(X^{j}_{\tau^{k}(i)})+W_{\tau^{k}(i)}.$ While the random vectors $X^{j\neq k}$ and $W$ have i.i.d entries, $X^{k}_{L}$ was obtained by sorting the coordinates of $X^{k}$. Thus, its coordinates are non-identically distributed and dependent. To solve the first problem, observe that the empirical mean is invariant under permutation and we can thus randomly shuffle the elements of $Y^{k}_{L}$ in order to obtain a vector with identically distributed coordinates. Furthermore, by De Finetti’s Theorem, any random vector with coordinates that are identically distributed (but not necessarily independent), can be viewed as a mixture of i.i.d vectors, effectively solving the second problem. Formally, the following result holds.

Lemma 6.3 shows that there are two sources of randomness in the distribution of $\widetilde{Y}^{k}_{L,i}$: the mixing variable $\Theta$ and the sub-Gaussian randomness of $\widetilde{X}^{k}_{L}|\Theta$ and $(X^{j\neq k},W)$. For the second source, it is possible to use standard concentration inequalities to show that conditioned on $\Theta$, $\widehat{\mathbb{E}}\left[\widetilde{Y}^{k}_{L}\right]$ concentrates around $\mathbb{E}\left[\widetilde{Y}^{k}_{L,1}|\Theta\right]$. We will formally do this in Lemma 6.5. Before we do this however, we focus on the first source and how $\Theta$ affects the distribution of $\widetilde{Y}^{k}_{L,i}$.

Since $\Theta$ is sampled from $\text{Beta}(n_{L}+1,n_{R})$, we can use standard techniques to show that it concentrates around $\frac{1}{2}$. More formally, we can use the following lemma, the proof of which is in the appendix.

In this section, we provide further justification of our theoretical results in the form of simulations on the finite sample count for active feature recovery under different regimes, as well as the predictive power of a single sub-tree as compared to the full tree. We additionally contrast DStump with the widely studied optimal Lasso .

We first validate the result of Theorem 5.1 and consider the linear design with $p=200$ and design matrix entries sampled i.i.d. from $U(0,1)$ with additive Gaussian noise $\mathcal{N}(0,.1)$. Concretely, we examine the optimal number of samples required to recover approximately $95\%$ of the active features $s$. This is achieved by conducting a binary search on the number of samples to find the minimal such value that recovers the desired fraction of the active feature set, averaged across 25 independent replications. In the leftmost plot of Figure 1, we plot the sample count as a function of varying sparsity levels $s\in[5,100]$ for DStump with a median split, DStump with the optimal split, as well as Lasso for bench-marking (with penalty parameter selected through standard cross-validation). By fixing $p$, we are evaluating the dependence of $n$ on the sparsity level $s$ alone. The results here validate the theoretical $\mathcal{O}(s\log p)$ bound that nearly matches the optimal Lasso . Also of note, the number of samples required by the median splitting is less than that of the optimal. Thus, in the linear setting, we see that DStump with median splitting is both more simplistic and computationally inexpensive. This optimal bound result is repeated with Gaussian data samples in the right most plot of Figure 1. Notably, in this setting the optimal split decision stumps perform better than the median as it demonstrates their varied utility under different problem contexts.

We additionally reiterate that the prior work of (Kazemitabar et al. 2017) attempted to simplify the analysis of the sparse recovery problem using DStump by examining only the left sub-tree, which produced the non-optimal $\mathcal{O}(s^{2}\log p)$ finite sample bound. To analyze the effect of this choice, the middle plot of Figure 1 presents the optimal sample recovery count when using only the left sub-tree subject to the additive model of Theorem 5.1. In accordance with our expectation and previous literature’s analysis, we see a clear quadratic relationship between $n$ and $s$ when fixing the feature count $p$.

Overall, these simulations further validate the practicality and predictive power of decisions stumps. Benchmarked against the optimal Lasso , we see a slight decrease in performance but a computational reduction and analytic simplification.

In this paper, we presented a simple and consistent feature selection algorithm in the regression case with single-depth decision trees, and derived the finite-sample performance guarantees in a high-dimensional sparse system. Our theoretical results demonstrate that this very simple class of weak learners is nearly optimal compared to the gold standard Lasso . We have provided strong theoretical evidence for the success of binary decision tree based methods in practice and provided a framework on which to extend the analysis of these structures to arbitrary height, a potential direction for future work.

The work is partially support by DARPA QuICC, NSF AF:Small #2218678, and NSF AF:Small # 2114269. Max Springer was supported by the National Science Foundation Graduate Research Fellowship Program under Grant No. DGE 1840340. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation.