Low-Rank Softmax Can Have Unargmaxable Classes in Theory but Rarely in Practice (2022)

Abstract: Classifiers in natural language processing (NLP) often have a large number of output classes. For example, neural language models (LMs) and machine translation (MT) models both predict tokens from a vocabulary of thousands. The Softmax output layer of these models typically receives as input a dense feature representation, which has much lower dimensionality than the output. In theory, the result is some words may be impossible to be predicted via argmax, irrespective of input features, and empirically, there is evidence this happens in small language models (Demeter et al., 2020). In this paper we ask whether it can happen in practical large language models and translation models. To do so, we develop algorithms to detect such unargmaxable tokens in public models. We find that 13 out of 150 models do indeed have such tokens; however, they are very infrequent and unlikely to impact model quality. We release our algorithms and code to the public.

## introduction probabilistic multiclass classifiers with a large number of output classes are commonplace in nlp @xcite . for example, the vocabulary size of contemporary lms and mt models varies from tens to hundreds of thousands @xcite . recent advances in modelling such large vocabularies have mostly been made by improving neural network feature encoders @xcite . but irrespective of a feature encoder's expressivity @xcite , a classifier that linearly maps lower dimensional features to higher dimensional outputs has reduced expressivity @xcite , with consequences that are not well understood. in this work we elaborate on the consequences of using argmax prediction with low-rank classifiers, on the left, each feature vector x is colored according to the class assigned the largest probability; note that while c 1 , c 2 and c 3 surface as regions, c 4 does not. on the right, we show that there is no direction in feature space for which c 4 has the largest probability. while @xcite proposed an algorithm to detect unargmaxable tokens and provided evidence of their existence in small lms, their proposed algorithm provided no guarantees and they were unable to test large lms. in this paper we ask: do unargmaxable tokens exist in large models used in practice? to answer this question, we develop algorithms to identify such tokens unambiguously. we tested 7 lms and 143 mt models. out of those, only 13 of the mt models exhibit unargmaxable tokens, and even for those cases the tokens are all noisy and infrequent. we conclude that although the expressivity constraints of lowrank softmax may have important ramifications, most practitioners do not need to worry about tokens that are unargmaxable. we provide new tools for them to confirm this on their own models. our contributions are the following: ## detecting unargmaxable classes in order to quantify whether unargmaxable classes arise in released lms and mt models, we first need to introduce tractable algorithms for detecting them. in this section we explain how unargmaxable classes can arise due to a softmax bottleneck. then, we introduce a fast approximate algorithm and a slow exact algorithm which we combine to detect vocabulary tokens that cannot be predicted. ## experiments in this section we use the combined algorithm from figure 4 to search models for unargmaxable tokens. we test 7 lms and 143 mt models. we find that unargmaxable tokens only occur in 13 mt models, but these are mostly infrequent and noisy vocabulary tokens. we therefore do not expect such tokens to affect translation quality per se. we also find that nearly all vocabulary tokens of lms and student mt models can be verified with less than @xmath0. in contrast, other mt models need thousands of steps and also rely on the exact algorithm. in this sense, models that need fewer steps are easier to verify: the search problem for their arrangement of softmax weights is easier. throughout the following experiments we assumed the softmax inputs were bounded in magnitude for all dimensions -100 ≤ x @xmath1. as we mentioned in subsection 3.2.1, if we have a softmax bias term, there are bounded regions. if the bounded regions are large, even though the outputs are not theoretically bounded, they are practically bounded since neural network feature encoders cannot produce arbitrarily large activations and some regions may be unreachable foot_8 . for the approximate algorithm, we search for a solution with a patience of @xmath2. we use gurobi (gurobi optimization, 2021) as the linear programme solver. we accessed the model parameters either via numpy @xcite or pytorch @xcite . the experiments took 3 days to run on an amd 3900x 12-core cpu using 10 threads and 64gb of ram. ## discussion we conclude from our experiments that it is possible to have unargmaxable tokens, but this rarely occurs in practice for tokens that would lead to irrecoverable errors in the mt models we checked. a limitation of our conclusions is that beam search is usually preferred over greedy decoding for mt models used in practice. we leave the question of whether unargmaxable tokens also impact beam search for future work. it is challenging to make exact claims about what can cause tokens to be unargmaxable because the models we tested varied in so many ways. however, we outline some general trends below. ## conclusions and future work in this work we discretised the outputs of softmax and showed how dimensionality constraints shrink the set of feasible class rankings and can lead to some classes being impossible to predict using argmax. in our experiments we demonstrated that while mt models can have unargmaxable vocabulary tokens, this does not occur often in our experiments. moreover, for the models we tested the unargmaxable tokens would not create discernible differences in translation quality as the tokens are noisy and infrequent. we release an algorithm to detect whether some classes are unargmaxable with the hope that this will be helpful to the wider community working on a plethora of different models where the observed phenomena may vary. in future work, we aim to investigate any learnability consequences more closely. as we saw, when using an approximate search algorithm, it is much harder to find argmaxable classes in some models than it is in others. since gradient descent algorithms are also iterative search algorithms seeking optimal parameters, we hypothesise that it will be challenging to train neural network encoders to map activations to regions of the input space that a search algorithm cannot find easily. hence, although some tokens may not be provably unargmaxable because of constraints imposed by the softmax parameters of the last layer, some tokens may still be very hard to produce because of difficulties encountered by the feature encoder. to this end, a more holistic investigation into the consequences of the loss in expressivity in low-rank classifiers is warranted.