SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jack H. Good, Torin Kovach, Kyle Miller, Artur Dubrawski
{jhgood,tkovach,mille856,awd}@andrew.cmu.edu
Decision trees are regarded for high interpretability arising from their hierarchical partitioning structure built on simple decision rules. However, in practice, this is not realized because axis-aligned partitioning of realistic data results in deep trees, and because ensemble methods are used to mitigate overfitting. Even then, model complexity and performance remain sensitive to transformation of the input, and extensive expert crafting of features from the raw data is common. We propose the first system to alternate sparse feature learning with differentiable decision tree construction to produce small, interpretable trees with good performance. It benchmarks favorably against conventional tree-based models and demonstrates several notions of interpretability of a model and its predictions.
In recent years, AI has made its way into increasingly critical applications, leading to the popularization of the topic of Trustworthy AI; AI should not only perform well in its primary task, but exhibit properties that earn trust, such as robustness, provable adherence to safety specifications, and interpretability to humans. For example, predictors that assist medical diagnosis are more useful and more likely to be trusted if their predictions are easily understood by medical personnel.
Interpretability is highly subjective and application-dependent by nature, and there is no universal standard of what makes a system highly interpretable. However, within machine learning (ML), there are some generally agreed-upon rules of thumb; for example, smaller models are considered more interpretable than structurally similar but larger ones, and rule-based representations such as decision trees are considered more interpretable than "black box" models such as neural networks.
Decision trees in particular are often used as the quintessential example of interpretable ML; they can be interpreted holistically (also called global interpretation) as a hierarchical partitioning of the input based on simple decision rules, and individual predictions can be interpreted (also called called local interpretation) as a series of such rules on a path from the tree's root to a leaf. Conventionally, these decision rules have the form  $x_{i} < t$  for some feature  $i$  of input vector  $\mathbf{x}$  and decision threshold  $t$ , thus representing axis-aligned hyperplanes in the input space, but many other decision rules are possible. For example, rules of the form  $\sum_{i}a_{i}x_{i} < b$  for some  $\mathbf{a}$ ,  $b$  are called oblique splits, and the trees using them oblique decision trees (ODTs).
In practice, however, there are barriers to the practical application of decision trees as interpretable models. First, due to various factors such as greedy growth, axis-aligned decision rules, and the lack of any objective to promote robust, large-margin predictions, they tend to generalize poorly to yet-unseen data; while this can be somewhat alleviated by pruning, the most effective and popular solution is to use large ensembles with strategies such as bagging as in Random Forests [9], random thresholds as in ExtraTrees [22], or boosting as in XGBoost [11]. In addition, decision trees often grow quite deep and use many decision rules for any given prediction, even for extremely simple data. Together, these make both global and local interpretation much more difficult.
(a) A decision tree is relatively complex and generalizes poorly.
(b) A random forest is very complex and generalizes better, but not perfectly.
(c) After rotation, a decision tree is simple and generalizes perfectly.
(d) Other data needs a different transformation to learn a simple, general tree.
Figure 1: An example of trees fitted to toy data to motivate the potential benefit of learning simple feature transforms. Points were sampled uniformly in the bands denoted by dashed lines.
We posit that these barriers are due, at least in part, to the sensitivity of decision trees to transformations of the input resulting from greedy construction and simple decision rules. Of these, key limitation is the latter; even if we replace greedy construction with a perfect tree learner, simple distributions can nonetheless require an arbitrarily large axis-aligned tree to fit. However, after a transformation of that distribution, even a greedy algorithm can produce a small tree, as shown in Section 3.5. Figure 1 gives a motivating example of a case where trees do poorly as interpretable learners on simple data, but succeed after a transformation. The sensitivity to the feature representation is also related to other limitations of tree-based models, such as their dependence on features hand-crafted using expert knowledge to achieve good performance on some problems, and their overall weakness on other problems, such as image classification.
This motivates the idea of learning transformations of the input features, or equivalently, more expressive decision rules such as oblique splits, to strike a balance between complexity of the rules and complexity of the overall model. This learning problem is not easy; axis-aligned decision rule learning, as in the popular CART algorithm [8], is only practical because the search is restricted to a finite set of candidate rules that is efficiently exhaustively searchable. However, finding even optimal oblique splits, which are among the simplest decision rules, is NP-Hard [26], and even oblique splits may not be considered interpretable unless they are sparse, that is, having few nonzero coefficients.
Nonetheless, methods have been developed to learn such trees, either by heuristically choosing splits in a tree construction process, or by constructing an axis-aligned tree and then optimizing the parameters of the splits (see Section 2 for examples). However, these suffer from either greedy, suboptimal choice of split, or from keeping the original structure of the tree. They are also usually restricted to a certain kind of split, such as oblique splits.
Here we propose a novel, flexible framework that is the first to alternate optimization of feature transformation parameters with complete regrowth of the tree. This is made possible by recent advances in fuzzy decision tree methods. Refitting the tree completely throughout training reshapes the tree and moves thresholds as needed; it also maintains the property that each splitting rule is maximally informative given the available features and data. Together these result in smaller, more interpretable trees for a given performance. Moreover, the flexibility of this framework to support various feature transformations provides a foundation for future work focusing on interpretable prediction with other data modalities, such as images and time series.
Decision trees are one of the oldest techniques in machine learning still widely in use today. It has long been known that constructing optimal binary decision trees is NP-Complete [29], so most algorithms use a greedy approach to "grow" a tree from root to leaf by successively partitioning the data by choosing the split that maximally improves the purity of labels in each partition. The partitions are usually axis-aligned, that is, involving only one feature of the input, to make the search practical. One of the most popular decision tree algorithms today is Classification And Regression Trees (CART) [8]. Unconstrained decision trees can always fit separable data perfectly, but are prone to overfitting; to improve generalization, growth can be constrained, but more often, trees are used in
randomized ensembles such as random forests [9] and ExtraTrees [22], or in boosted ensembles such as AdaBoost [19] and XGboost [11].
There is extensive work using decision trees extracted from more complex models either to entirely replace the model, or as an interpretable proxy. Examples include [12, 37, 43, 3, 35, 36, 25, 7].
Other work attempts to improve the interpretability of decision trees themselves by introducing new decision rules. These rules are more complex than the axis-aligned threshold rules, but compensate by improving performance of single-tree models while reducing their size. For example, [6] uses a branch-and-bound algorithm to search for bivariate splits. [16, 17, 23] use parallel coordinates, which incorporate information from multiple data features while being easy to visualize. [13] uses bilevel optimization to find nonlinear splits for growing trees.
The most common augmented decision rule, however, is linear splits, that is, splitting using arbitrary hyperplanes rather than axis-aligned hyperplanes. These are often called oblique splits, and the trees that use them oblique decision trees (ODTs). Algorithms for growing ODTs are first proposed along with CART [8]. [34] later discusses the difficulty of finding optimal linear splits due to local minima and proposes a heuristic using randomization and deterministic hill climbing.
Another approach is to grow an axis-aligned tree, fix the structure, then treat it as a parametric model in order to learn coefficients for linear splits. This can be made possible using fuzzy splits, wherein a decision is a weighted combination of both subtrees rather than wholly one or the other, to make the tree differentiable. These have gone by names such as fuzzy, soft, neural, or differentiable decision trees; we use fuzzy decision trees (FDTs) as an umbrella term. In the case of ODTs, a splitting function such as sigmoid is used to map a linear combination of feature values to the weight of each subtree. [42] is the first to train globally optimized trees in this way. Later, [32] proposes ensembles of such trees, both as standalone models and as the final layer of a neural architecture, [33] presents various improvements such as balance and sparsity regularizations and dynamic learning rate adjustments, and [40] applies them to reinforcement learning.
Tree Alternating Optimization (TAO) is a recent and successful ODT learning algorithm that initializes similarly, but unlike FDTs, uses a linear discriminator at each node, such as logistic regression or linear SVM, then alternates optimization over each depth level of the tree. It can be used with a sparsity penalty to produce sparse ODTs, that is, where many or most of the coefficients are 0, making linear splits more interpretable. While it is shown that this alternating optimization never causes the loss to increase, there are no convergence or approximation guarantees. It was initially proposed for classification [10], then extended in various ways, including regression ensembles [45]; boosted ensembles [21]; clustering [20]; semi-supervised learning [46]; interpretable image classification [41]; interpretable natural language processing [28]; and as a tool for understanding parts of neural networks [27].
Our approach is different from any existing work in that we consider a broader scope of feature types, but more importantly, in that we continually and completely refit the decision tree throughout the learning process; this requires an efficient differentiable decision tree. For this, we rely on Kernel Density Decision Trees (KDDTs) [24], a recently proposed FDT that represents uncertainty through kernels rather than splitting functions and learns axis-aligned fuzzy splits using a generalization of CART. The kernel representation also has the benefit of avoiding redundant splitting that may occur in other FDT formalisms because repeating similar splits can make the division less "fuzzy" and therefore reduce loss by inflating the tree size. See the paper for background on axis-aligned FDTs. Our approach is also unique in that, for each output of the feature transformation, we weight the regularization by how much the feature is used for decisions in the training data, resulting in a more balanced enforcement of sparsity throughout the tree.
3.1 Alternating optimization trees and features
Given training data with features  $\mathbf{X} \in \mathbb{R}^{n \times p}$  and regression targets or one-hot class label vectors  $\mathbf{Y} \in \mathbb{R}^{n \times q}$ , we aim to learn parameters  $\theta \in \mathbb{R}^d$  of feature transformation  $f_{\theta}: \mathbb{R}^p \to \mathbb{R}^{p'}$  and a decision tree  $T: \mathbb{R}^{p'} \to \mathbb{R}^q$  to minimize a loss function  $\mathcal{L}(\theta, T, \mathbf{X}, \mathbf{Y})$  for the composed predictive
model  $T \circ f_{\theta}$ . The tree is typically subject to constraints such as a max depth, minimum sample weight per node, or cost-complexity pruning.
Viewed within this framework, the prevailing approach in prior work is equivalent to the following process: (1) initialize  $\theta$  such that  $f_{\theta}$  is identity; (2) with  $\theta$  fixed, optimize  $T$  using a greedy tree growth strategy such as CART; (3) with  $T$  fixed, optimize  $\theta$  using gradient-based optimization. Some methods use a growing and/or pruning heuristic during (3) so that the structure of  $T$  may change.
Instead, we propose an approach inspired by alternating optimization, a common strategy to pragmatically solve difficult optimization problems, whereby we alternate fitting a tree as in (2) with one or more gradient-based updates to  $\theta$  as in (3). In this case, alternating allows us to use different optimization strategies for the feature transform parameters  $\theta$  and the non-parametric decision tree  $T$ . At the end, we can use the fuzzy tree  $T$ , or replace its fuzzy splits with conventional crisp splits; the latter improves interpretability when predictions would otherwise take multiple paths, but can result in a small drop in performance.
3.2 Efficient differentiable trees
Conventional decision trees are piecewise-constant, so there is no meaningful loss gradient with respect to  $\theta$ . Instead, we use kernel density decision trees (KDDTs) [24], a type of fuzzy decision tree that interprets inputs with local uncertainty defined by a kernel function. In essence, they fit using a kernel density estimate computed from  $\mathbf{X}$  rather than the raw values, and predictions are smoothed over the same kernel. This smoothing makes both the predictions and loss differentiable. KDDTs are fitted using an efficient generalization of CART, making them suitable for applications such as this where frequently repeated fitting makes efficiency crucial.
Here we present a simplified formulation of KDDTs to focus on this use case. A KDDT  $T: \mathbb{R}^p \to \mathbb{R}^q$  is defined by underlying decision tree  $T_0: \mathbb{R}^p \to \mathbb{R}^q$  and product kernel  $k: \mathbb{R}^p \to \mathbb{R}$  with  $k(\mathbf{x}) = \prod_{i=1}^{p} F_i'(x_i)$ , and  $F_i$  the marginal CDF of the kernel along each dimension  $i$ . For the KDDT fitting algorithm, each  $F_i'$  must be piecewise-constant, but this restriction does not apply for prediction; for example, one might use a histogram approximation of a Gaussian distribution for fitting and the Gaussian itself for prediction.
Then the KDDT is the output element-wise convolution  $T^{(j)} \equiv T_0^{(j)} * k$ . Suppose  $T_0$  partitions  $\mathbb{R}^p$  into hyper-interval leaves  $L \subseteq \mathbb{R}^p$  with bounds  $\ell_{L,i}$ ,  $u_{L,i}$  (which may be  $-\infty$  or  $\infty$ ) on each dimension  $i$ . Each leaf has a value  $\mathbf{v}_L \in \mathbb{R}^q$ . Then the prediction is
$$
T ^ {(j)} (\mathbf {x}) = \left(T _ {0} ^ {(j)} * k\right) (\mathbf {x}) = \sum_ {L} \mu_ {L} (\mathbf {x}) v _ {L} ^ {(j)} \tag {1}
$$
where  $\mu_{L}$  is the leaf membership function  $\mu_{L}(\mathbf{x}) = \prod_{i\in [p]}F_{i}(x_{i} - \ell_{L,i}) - F_{i}(x_{i} - u_{L,i})$  and  $\mathbf{v}_L\in \mathbb{R}^q$  is the leaf value  $v_{L}^{(j)} = \frac{\sum_{i\in[n]}\mu_{L}(\mathbf{X}_{i})Y_{i,j}}{\sum_{i\in[n]}\mu_{L}(\mathbf{X}_{i})}$ . The Jacobian is thus
$$
\nabla_ {i, j} T (\mathbf {x}) = \sum_ {L} \nabla_ {i} \mu_ {L} (\mathbf {x}) v _ {L} ^ {(j)}
$$
$$
\nabla_ {i} \mu_ {L} (\mathbf {x}) = (F _ {i} ^ {\prime} (x _ {i} - \ell_ {L, i}) - F _ {i} ^ {\prime} (x _ {i} - u _ {L, i})) \prod_ {i ^ {\prime} \neq i} F _ {i} (x _ {i} - \ell_ {L, i}) - F _ {i} (x _ {i} - u _ {L, i})
$$
using the same traversal algorithm to compute (1) from [24]. Trees are fitted using an information gain based on an impurity measure, so the same impurity should be used as the loss function for feature learning. For classification, we use the Gini index; the resulting loss function, including the feature transformation  $f$ , is  $\mathcal{L}(\theta, T, \mathbf{X}, \mathbf{Y}) = \frac{1}{n} \sum_{L} w_{L} (1 - \mathbf{v}_{L}^{T} \mathbf{v}_{L})$  with  $w_{L} = \sum_{i \in [n]} \mu_{L}(f_{\theta}(\mathbf{X}_{i}))$  the sample weight at leaf  $L$ . Then we compute the loss gradient  $\nabla_{\theta} \mathcal{L}$  in terms of  $\mu, \nabla \mu$ , and  $\nabla_{\theta} f$ , using the sparsity of  $\mu$  for efficient computation during tree traversal. For regression, we use variance of regression target values as the impurity, which can be expressed as a shifted version of the Gini loss, so the gradient computation is the same.
3.3 Kernel choice
In this context, the kernel is mainly used as a smoother to make the tree differentiable, so the choice of shape is not crucial. Good choices are a uniform (box) or Gaussian kernel. However, the size of
the kernel relative to the scale of the input is important; we adopt the term bandwidth from kernel density estimation to describe kernel scaling. If the bandwidth is too small, a KDDT is very close to a conventional decision tree, and the loss gradient is either zero or highly unstable because it results from very few data near split boundaries; if the bandwidth is too large, however, splits are very soft, the tree grows large, the discriminative power is weak, and computation is slow due to dense membership of data in leaves. To control the bandwidth, we use one of two strategies: (1) design the feature transform so that the output range is limited, for example, to [0, 1], and choose an appropriate kernel bandwidth; or (2) use a regularization on the feature transform parameters to automatically scale the feature outputs appropriately relative to the kernel. We generally find the latter approach to be more effective and describe the use of regularizations in detail in Section 3.5.
Another related consideration is the support of the kernel, that is, the size of the range on which it has nonzero density. In general, the larger the support, the more decision paths that data belong to, increasing the cost of fitting, gradient computation, and prediction (if using a kernel during prediction). For example, a Gaussian kernel truncated at 3 standard deviations and another with the same variance but truncated at 6 standard deviations are nearly identical, but the latter is more computationally expensive to use with a KDDT.
3.4 Parameterized feature transforms
In principle, any differentiable parameterized class of feature transforms can be used, ranging from simple linear transforms as in ODTs to expressive classes of functions such as MLPs that have high discriminative power on their own. Generally, more expressive feature transforms result in smaller trees, but are probably less interpretable, depending on the application. The goal then is to choose expressive feature classes that remain interpretable in the desired context when conjoined into a rule list, as in the decision path of a tree.
In this work, we focus primarily on tabular data, to which tree-based models are most successfully applied, and suggest some general feature transformation primitives based on feature engineering techniques typically used with tree-based models. These can be composed, or their outputs concatenated so different kinds of transforms can be used together in the tree. We generally recommend starting with linear and prototype features, then proceeding with whatever improves performance while meeting interpretability requirements for the application.
Before describing the primitives, we note a general consideration for features used with decision trees: decision trees consider only the order of inputs and not the distance between values. While this is not strictly true for fuzzy decision trees such as KDDTs, it is best to avoid adding unnecessary feature complexity by using element-wise monotonic transformations of features, such as bias terms, monotonic activation functions, etc., on the final output of the feature transformation.
Identity. It may be beneficial to include the original, unmodified features along with transformed features so that the simplest possible rules are considered during construction of the tree. This also, in some sense, safeguards against situations where transformed features may actually be less informative than the original features.
Element-wise transformation. As noted above, element-wise monotonic transformations are largely meaningless on their own, but can be useful if composed with other transforms. For example, the composition of element-wise square with linear transformation allows for conic decision rules. Some element-wise mappings worth considering have no parameters, such as  $x \mapsto \exp(x)$  and  $x \mapsto \log(x)$ ; others have parameters to be learned, such as  $x \mapsto x^a$ . A notable option is the element-wise shifted power transformation
$$
x \mapsto \left\{ \begin{array}{l l} \frac {\operatorname {s g n} (x + \alpha) | x + \alpha | ^ {\lambda} - 1}{\lambda (x + \alpha) ^ {\lambda - 1}} & \text {i f} \lambda \neq 0 \\ (x + \alpha) \operatorname {s g n} (x + \alpha) \log (x + \alpha) & \text {i f} \lambda = 0 \end{array} \right.
$$
which is used to alter distribution shape, typically to make values more normally distributed.
Linear transformation. The linear transformation is  $\mathbf{x} \mapsto \mathbf{A}\mathbf{x} + \mathbf{b}$  with weights  $\mathbf{A}$  and bias  $\mathbf{b}$ . The bias is used only if composing with another primitive. Using only a linear transformation results in an ODT. These can be initialized randomly, as identity, or using a linear dimension reduction strategy such as PCA.
data
n, p, q	LR	MLP	DT	RF	ET	XGB	ours:
fuzzy	linear
crisp	ours:
fuzzy	proto
crisp
iris [18]	0.960	0.953	0.947	0.947	0.953	0.947	0.960	0.960	0.947	0.927
150, 4 (4), 3	-	-	6.4	7.2e2	2.1e3	4.3e2	6.1	7.6	6.5	2.9
heart-disease [30]	0.822	0.792	0.707	0.802	0.795	0.792	0.812	0.812	0.793	0.779
303, 13 (20), 2	-	-	13.9	4.8e3	1.1e4	7.9e2	21.6	19.4	5.0	3.9
dry-bean [31]	0.925	0.934	0.912	0.923	0.921	0.928	0.920	0.913	0.915	0.901
13611, 16 (16), 7	-	-	99.8	6.7e4	2.0e5	1.3e4	1.1e2	45.8	1.6e2	94.2
wine [1]	0.983	0.989	0.904	0.977	0.989	0.955	0.983	0.983	0.961	0.961
178, 13 (13), 3	-	-	8.5	9.4e2	3.3e3	2.4e2	2.0	2.0	2.3	2.3
car [5]	0.926	0.992	0.977	0.964	0.971	0.994	0.991	0.992	0.980	0.979
1728, 6 (21), 4	-	-	95.3	2.3e4	3.1e4	4.5e3	29.0	29.0	59.5	55.9
wdbc [44]	0.974	0.975	0.935	0.965	0.970	0.968	0.972	0.972	0.961	0.961
569, 30 (30), 2	-	-	13.0	1.9e3	6.0e3	2.7e2	1.3	1.3	6.4	5.0
sonar [38]	0.755	0.879	0.735	0.826	0.880	0.855	0.818	0.799	0.798	0.817
208, 60 (60), 2	-	-	14.1	2.0e3	5.6e3	3.0e2	5.7	3.9	10.3	10.3
pendigits [2]	0.952	0.994	0.964	0.993	0.994	0.991	0.981	0.976	0.950	0.931
10992, 16 (16), 10	-	-	3.2e2	3.8e4	9.8e4	8.5e3	2.6e2	2.4e2	3.1e2	3.1e2
ionosphere [39]	0.875	0.917	0.892	0.934	0.943	0.943	0.932	0.920	0.920	0.909
351, 34 (34), 2	-	-	15.5	2.2e3	5.9e3	3.4e2	3.9	5.5	10.0	6.3
Table 1: Results of tabular data benchmarks. Number of attributes  $p$  is listed before and after one-hot encoding categorical attributes. Each cell shows 10-fold cross-validation accuracy and average number of splits. The best accuracy for each data set is in bold.