Speeding Up Entmax (2022)

Abstract: Softmax is the de facto standard for normalizing logits in modern neural networks for language processing. However, by producing a dense probability distribution each token in the vocabulary has a nonzero chance of being selected at each generation step, leading to a variety of reported problems in text generation. \alpha -entmax of Peters et al. (2019) solves this problem, but is unfortunately slower than softmax. In this paper, we propose an alternative to \alpha -entmax, which keeps its virtuous characteristics, but is as fast as optimized softmax and achieves on par or better performance in machine translation task.

## introduction sparseness of vector representations is a desirable trait in neural network models for natural language processing (nlp): words (subwords) are discrete objects by their nature, and, accordingly, are encoded by one-hot embeddings at the input and output of neural networks. however, to predict a categorical response in neural models, softmax is most often used, which produces a dense probability distribution, i.e. every category (word/subword) receives a non-zero probability. recent studies suggest that it is this output density that poses problems when the trained nlp model is used for inference. for example, in the case of text generation, unconstrained sampling from a trained language model results in poor quality of the resulting text @xcite . in neural machine translation (nmt), exact decoding from a trained model often results in empty text @xcite . 1 to get around these problems, constrained decoding techniques have been proposed, most of which artificially impose sparsity on softmax prediction. for example, @xcite propose to sample from the top-k probable words, and @xcite propose to sample from the most probable words, which comprise the cumulative probability p. while these methods are effective, they are ad-hoc solutions that lead to a mismatch between how the model is trained and how it is used at inference. in this regard, the works on sparse alternatives to softmax stand apart since they allow us to make inference from the model in the same way than it was trained. some of the most successful and elegant solutions are sparsemax @xcite and its generalization α-entmax @xcite . when coupled with suitable losses, these transformations are not inferior to softmax, and sometimes even surpass it as measured with final performance metrics on a number of tasks. a problem with these transformations however is that they are significantly slower than softmax when the number of categories (vocabulary size) is tens of thousands, as in the case of text generation. this is because α-entmax transformation-in its original formulation-requires sorting over the logits. 2 in this work, we ask the question: is it possible to obtain a sparse output like that of α-entmax, but without its degradation in computational speed? our answer is affirmative-we propose a sparse output transformation that • is on par or superior to softmax and α-entmax in the nmt tasks, the most surprising thing is that such a transformation is simply a shifted relu raised to power 1 α-1 , which we call α-relu. the rest of the paper is organised as follows. in sect. 2 we motivate the choice of α-relu as the output transformation, and also select an appropriate loss function. in sect. 3 we experimentally confirm our claims about performance and output speed of α-relu in the nmt task. sect. 4 is devoted to a comparative analysis of α-relu and α-entmax in terms of sparsity, ability to solve the empty translation problem, and training dynamics. ## α-relu at output our departure point is the α-entmax transformation of @xcite which can be defined for @xmath0 as we can see, the threshold τ is only needed to ensure that α-entmax(z) is a probability distribution. we loosen this constraint, and only require non-negative weights, which is sufficient for most uses. consider then a transformation proof. the proof is in appendix b.1. by lemma 1, gradient-based minimization of ℓ indeed forces α-relu(z) → e y . notice that this is similar to what happens when the softmax normalization is coupled with the cross-entropy loss or when α-entmax is coupled with the entmax loss. in both cases differentiating the loss with respect to logits gives p-e y , where p is either softmax(z) or α-entmax(z) @xcite . ## experiments in theory, nothing prevents α-relu from learning what α-entmax is learning. however, in practice we can have a different picture, because training is conditioned by many factors-the size of the dataset, the architecture of the neural network, the optimization algorithm, etc. in this section, we compare α-relu empirically with α-entmax (as well as with sparsemax and softmax), assuming all other factors are fixed. the goal of these experiments is to evaluate the consequences of using α-relu as drop-in replacement for α-entmax. ## limitations and risks we believe that the main limitations of our work are as follows: the last mentioned limitation leads to the potential risk of inability to learn if α-relu is misused in the intermediate layers of the neural network such as attention layers. the experiments of @xcite using vanilla relu (2-relu with τ = 0 in our notation) instead of softmax to produce attention weights lead to a divergence of the loss function of the transformer model. this translates into a waste of energy, especially when training large models on large datasets. therefore, we believe that in the future, a preliminary mathematical analysis and/or experiments with small models on small datasets should be carried out as to why the unnormalized distribution of attention weights leads to the inability of the model to learn. ## conclusion it seems that the sparsity of the output is natural for (sub)word prediction models. nevertheless, sparsity does not have to come with slowdown of computations, as our work shows. the proposed transformation, α-relu, gives a sparse output, shows competitive performance, and is as fast as softmax. the reduced dependency on the vocabulary size seems particularly important in translation, where neural models are moving more and more towards multi-lingual ones, which in general have a much higher vocabulary size in order to accommodate enough tokens for all languages. a natural extension of this work will be the evaluation of α-relu in the problem of open-ended text generation, as well as a replacement for softmax in the attention layers of transformer models. our standalone implementation of α-relu in pytorch is available at https://github. com/maxattezekbayev/alpha-relu .