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

Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10730–10745 November 7–11, 2021. c⃝2021 Association for Computational Linguistics 10730 Softmax Tree: An Accurate, Fast Classifier When the Number of Classes Is Large Arman Zharmagambetov Magzhan Gabidolla Miguel Á. Carreira-Perpiñán Dept. of Computer Science and Engineering, University of California, Merced, USA {azharmagambetov,mgabidolla,mcarreira-perpinan}@ucmerced.edu Abstract Classification problems having thousands or more classes naturally occur in NLP, for ex- ample language models or document classifi- cation. A softmax or one-vs-all classifier natu- rally 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 bi- nary tree having sparse hyperplanes at the de- cision nodes (which make hard, not soft, de- cisions) and small softmax classifiers at the leaves. This is much faster at inference be- cause 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 optimiza- tion (TAO). Compared to a softmax and other classifiers, the resulting softmax trees are both more accurate in prediction and faster in infer- ence, as shown in NLP problems having from one thousand to one hundred thousand classes. 1 Introduction Classification problems with thousands or more classes (sometimes called extreme classification) naturally occur in NLP and other areas. One ex- ample are language models. There are about 171k words in the current edition of the Oxford En- glish 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 on- tology scheme. In this many-class setting, it is con- siderably 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) clas- sifier, 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 in- stance, 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 num- ber of classes, which is slow when K and D are large. This problem also occurs with other clas- sifiers, such as soft decision trees. Indeed, com- putational constraints on the vocabulary size are a major challenge for neural machine translation (Koehn, 2020), for example. We argue that having the classifier output a pos- itive 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 …

Content truncated to 500 words