Softmax Tree: An Accurate, Fast Classifier When the Number of Classes Is Large (2021)

Abstract: Classification problems having thousands or more classes naturally occur in NLP, for example language models or document classification. A softmax or one-vs-all classifier naturally handles many classes, but it is very slow at inference time, because every class score must be calculated to find the top class. We propose the “softmax tree”, consisting of a binary tree having sparse hyperplanes at the decision nodes (which make hard, not soft, decisions) and small softmax classifiers at the leaves. This is much faster at inference because the input instance follows a single path to a leaf (whose length is logarithmic on the number of leaves) and the softmax classifier at each leaf operates on a small subset of the classes. Although learning accurate tree-based models has proven difficult in the past, we are able to overcome this by using a variation of a recent algorithm, tree alternating optimization (TAO). Compared to a softmax and other classifiers, the resulting softmax trees are both more accurate in prediction and faster in inference, as shown in NLP problems having from one thousand to one hundred thousand classes.

## introduction classification problems with thousands or more classes (sometimes called extreme classification) naturally occur in nlp and other areas. one example are language models. there are about 171k words in the current edition of the oxford english dictionary, and many more if we include all forms of a word, names, technical acronyms, etc. another example is document classification. the open directory project (odp) contains over 1m website categories organized in a hierarchical ontology scheme. in this many-class setting, it is considerably difficult to learn a model that is accurate and fast at inference time. the simplest and most widespread model is a linear (e.g. softmax) classifier, possibly as the output layer of a neural net. one important problem with a softmax classifier is that one must compute the score or probability of (nearly) all classes, conditional on the input instance, in order to determine the (top-n) predicted class. this has a cost o(dk) where d is the input dimension of the softmax and k the number of classes, which is slow when k and d are large. this problem also occurs with other classifiers, such as soft decision trees. indeed, computational constraints on the vocabulary size are a major challenge for neural machine translation @xcite , for example. we argue that having the classifier output a positive probability (however small) for each class is slow and unnecessary when k is large, because, for any given instance, the majority of classes should indeed have a negligible probability. a much faster classifier is a traditional decision tree, which makes hard decisions by thresholding an input feature at the decision nodes and outputs a single class at each leaf. this axis-aligned tree assigns zero probability to all classes except the predicted one, which is reached through a single rootleaf path very quickly (in log k time if the tree is balanced). however, such trees are known to be insufficiently accurate even if grown very deep. we propose a softmax tree (st), a binary tree having sparse hyperplanes at the decision nodes (which make hard, not soft, decisions) and a small softmax classifier outputting @xmath0. a st is still very fast at inference: it sends the input instance to a single leaf via a path whose length is logarithmic on the number of leaves (for a complete tree), and it assigns (without computing them) probability zero to most classes (namely, all classes not in the leaf). trading off the depth ∆ of the tree and the number of classes k per leaf can potentially result in fast, highly accurate classifiers. however, sts are still hard to train because they define a nonconvex, nondifferentiable prob-lem. we solve this by modifying tree alternating optimization (tao), a recent algorithm for learning oblique decision trees (having hyperplane decision nodes and constant-label leaves), so that it can handle softmax leaves, and by using a good initialization. before describing the training algorithm (section 4), we review related work (section 2) and show (section 3) that oblique tree classifiers with constant-label leaves are intrinsically more powerful than linear classifiers, but possibly less efficient, which justifies our sts as hybrid treesoftmax classifiers. finally, (section 5) we convincingly show that our sts have both higher accuracy and much faster inference time than several previous models on high-dimensional problems having up to 10 5 classes. ## related work the extreme multi-class classification problems have been previously addressed both in the literature of machine learning and natural language processing. the most basic method is one-versus-all @xcite where an independent binary classifier is learned per class. another classical approach, error correcting output codes (ecoc) by @xcite , represents each class with a binary code and learns a separate binary classifier for each bit. however, these methods become costly or even intractable (both for training and predicting) when the number of classes is huge. moreover, each binary classifier needs to handle highly imbalanced dataset as all classes except one would be instances of the negative class. various approaches have been proposed to speed up the training/prediction time and reduce the computational complexity. the most popular of them capitalize on using tree-based structures since it naturally leads to the logarithmic time reduction. decision trees have been actively used in this area @xcite . nevertheless, traditional axis-aligned decision trees, such as c4.5 @xcite or cart @xcite , have very low accuracy @xcite . nested dichotomies @xcite ) rely on a tree structure to divide a set of classes into two disjoint subsets and learn a binary classifier to separate them. however, human expertise is necessary to obtain a tree structure and class assignments. additionally, the total error of the model accumulates over the depths since there is no way to refine binary classifiers once split is performed. more recent works which are specifically designed to cope with large number of classes @xcite employ similar idea but take into account class distributions to generate a tree structure. other tree-based approaches include global or partial optimization over parameters of a tree. for instance, @xcite propose to use a fixed structured tree where each node has much smaller sized linear multi-class classifier. @xcite extend this work by allowing a tree structure to grow. other works capitalize on generating "perfectly" balanced trees to guarantee logarithmic inference time @xcite . optimizing a tree parameters in these methods is typically done by approximating gradient information in a certain way (possibly in "online" fashion). other tree-based methods exist with more focus on large scale extreme multi-label classification and ranking @xcite . finally, there are works which employ non-tree based approaches, such as hashingbased methods @xcite , subsampling of classes and training set @xcite , etc. in the context of nlp, most of the abovementioned methods are applicable in the number of practical applications, such as large-scale document classification and language modeling. moreover, there are methods that are specifically designed for language modeling tasks where vocabulary size can be very large and it demands efficient computation of the softmax outputs. hierarchical softmax (hsm) @xcite is an approximation which employs a "soft" decision tree with linear nodes to address this issue. hsm has been actively used in the problem of learning distributed representations of words @xcite where it can be jointly trained with neural nets of various complexity @xcite . follow up works on this topic @xcite propose various initializations for the tree structure (e.g. random, huffmann tree, etc.). although the training of hsm can be efficiently done using specific loss functions, but during prediction time, input follows all children with a certain probability which brings no speedup compared to the plain softmax. it is still possible to transform a soft tree back into a "hard" tree once training is done (by choosing a child with the highest probability at each split). for example, a recent work from @xcite apply a similar approach. however, as we will experimentally show later, it increases an error due to further approximation. similar observations were found in @xcite where various subsampling techniques showed better results. recently, certain pruning mechanisms have been proposed as an alternative approach to speed up the prediction time @xcite . the parameters are typically learned from data (e.g. by using the one-vs-all scheme or by optimizing a loss such as the cross-entropy for a softmax). theorem 3.1. any k-class linear classifier can be exactly represented by an oblique classification tree with constant-label leaves. the converse is not true. proof. we give a constructive proof. define @xmath1, so the linear classifier output is f (x) = arg max(z 1 , . . . , z k ). the latter argmax function of k arguments can be exactly computed by a complete binary tree of depth k -1 as illustrated in fig. 1 for @xmath2. each decision node performs a comparison "z @xmath3, else the left child. each leaf's class label is the corresponding argmax value. this tree represents the iterative algorithm to compute arg max(z 1 , . . . , z k ) by scanning each element left to right for the maximum: max(z k , . . . max(z 3 , max(z 2 , z 1 ))). specifically, each root-leaf path tree corresponds to one possible execution path, and the nodes at depth j compare with z j . since a comparison "z @xmath4, each decision node is a linear decision function and the tree is oblique. this argument is valid for interior points, but can be made to work for points on the boundary between classes by appropriately breaking the ties (replacing "≥" with ">") as needed. the converse is not true because an oblique tree can dedicate more than one leaf to a class, resulting in a nonconvex class region (the union of two polytopes). this cannot be represented by a linear classifier, whose class regions are polytopes of the form a i x + b @xmath5. note that an axis-aligned decision tree (where each decision node tests a single input feature), however deep, cannot exactly represent a linear classifier unless the latter is itself axis-aligned. the above theorem shows that oblique decision trees are strictly a more powerful family of classifiers than linear classifiers foot_0 . the proof construction produces a large tree, having 2 k-1 leaves (although its inference time is equal to that of the linear classifier, o(dk)). in practice, the tree need not exactly represent the linear classifier throughout the input space but just over the instances of a training set, and this can likely be achieved with far smaller trees. for example, if each class is linearly separable from the rest, the corresponding tree has just k leaves and each of the k -1 decision nodes has exactly one leaf child (except the deepest decision node, which has two leaf children). still, for a given dataset, the linear classifier may be a more efficient model in number of parameters than the tree. which classifier (linear or oblique tree with constant-label leaves) is better is an empirical question. a further issue is that finding the global optimum of the training problem is easy for a linear classifier (e.g. the cross-entropy is convex for a softmax classifier) but np-hard for a decision tree. this leads us to the model we propose in this paper, the softmax tree, which is a hybrid between a purely linear classifier and an oblique tree with constant-label leaves. ## softmax trees: definition and training unlike soft decision trees, which can be readily optimized via gradient-based methods, hard decision trees pose a far more difficult optimization problem, not just nonconvex but nondifferentiable (and np-hard). traditional tree learning algorithms such as cart @xcite or c5.0 @xcite are based on greedily and recursively partitioning the input space, and pruning the resulting tree to reduce overfitting. however, they are known to produce suboptimal trees @xcite . in this work, we build on a recent algorithm, tree alternating optimization (tao) @xcite which is a non-greedy optimization method for treebased models. originally described for oblique trees with constant-label leaves, tao has shown a huge success in training a wide range of other tree-based models: regression trees (zharmagambetov and carreira-perpiñán, 2020), tree ensembles @xcite , hybrid models @xcite , etc. moreover, they are shown to have a great potential to study model interpretability and explainability @xcite . tao works very differently from cart and much more like a regular machine-learning optimization algorithm, but instead of gradients (which do not apply) it uses alternating optimization over groups of nodes of a fixed tree structure. this results in a monotonic decrease of the objective function over all the tree parameters and convergence to a local optimum. detailed comparison of tao against traditional trees can be found in @xcite . next, we describe our softmax trees and training algorithm, noting the differences with carreira-perpiñán and tavallali (2018). consider a k-class problem with training set {(x n , y n )} n @xmath6, . . . , k} of ddimensional instances and labels. let t(x; θ) be a binary decision tree which produces a prediction for each input x by routing x from the root to exactly one leaf and applying a predictor function at that leaf. each node (both decision and a leaf) has learnable parameters θ i and the total set of parameters of a tree is θ = {θ i } @xmath7, where n is the set of nodes. each decision node i has a decision function f i (x; θ i ): r @xmath8, right i } ⊂ n , sending instance x to the corresponding child of node i, and each leaf has a predictor function g i (x; θ i ): r @xmath9, . . . , k} that produces the actual output. in a softmax tree (st): tao assumes a fixed tree structure (say, complete of depth ∆) and initial node parameters. hence, the hyperparameters of a st are ∆ and k. tao optimizes the following objective function: tao is based on two theorems. first, eq. ( 1 ) separates over any subset of non-descendant nodes (e.g. all the nodes at the same depth); this follows from the fact that the tree makes hard decisions. all such nodes may be optimized in parallel. second, optimizing over the parameters of a single node i simplifies to a well-defined reduced problem over the instances that currently reach node i (the reduced set r i ⊂ {1, . . . , n }). the form of the reduced problem depends on the type of node: this is as in carreira-perpiñán and tavallali (2018) except that in our sts the loss is the cross-entropy of the corresponding leaf. this problem is np-hard but can be well approximated with a convex surrogate; we use ℓ 1regularized logistic regression where each instance is weighted by the loss difference between the winner child and the other child, and solve it using liblinear @xcite . in our sts, g i is a k-class softmax classifier with an ℓ 1 sparsity penalty. we first estimate the k classes (out of k possible classes) as the k most populous classes in r i . then we train the softmax, which is a convex problem. we solve it using sag @xcite . the resulting algorithm visits nodes in reverse breadt-first search order and is shown in algorithm 1. essentially, each iteration trains all nodes at the same depth (in parallel) from the leaves to the root, by solving either an ℓ 1 -regularized softmax classifier at each leaf, or an ℓ 1 -regularized logistic regression at each decision node. note that the ℓ 1 penalty on the decision nodes' weight vectors means that some of them may become zero, which makes the node redundant and can be pruned at the end, reducing the size of the tree. ## experiments we demonstrate the performance of our method on two popular nlp tasks: (a) large scale text classification, and (b) language modeling. experiments suggest that our resulting softmax trees outperform simple and advanced baselines either in accuracy (and yet very fast) or in prediction time (and yet showing competitive accuracy); or quite often in both of these indicators. moreover, the resulting models are compact in terms of memory requirements. ## conclusion softmax trees strike a balance between having a single softmax classifier, which is easy to optimize but slow at inference, and a decision tree with contant-label leaves, which is hard to optimize but fast at inference. tuning the depth of the tree and the number of classes per leaf softmax results in classifiers that are both more accurate and much faster than a regular softmax or other hierarchical softmax approaches in many-class problems. finding good local optima for softmax trees is possible with a modification of the tree alternating optimization (tao) algorithm combined with a good initialization. we are now working on forests of softmax trees and on growing the tree structure adaptively.