SlowGuess/NAACL_2022_OCR huggingface dataset:
Maxat Tezekbayev<sup>1</sup>, Vassilina Nikoulina<sup>2</sup>, Matthias Galle<sup>2</sup>, and Zhenisbek Assylbekov<sup>1</sup>
$^{1}$ Nazarbayev University {maxat.tezekbayev, zhassylbekov}@nu.edu.kz  
 $^{2}$ NAVER LABS Europe {vassilina.nikoulina, matthias.galle}@naverlabs.com
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.
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 (Holtzman et al., 2020). In neural machine translation (NMT), exact decoding from a trained model often results in empty text (Stahlberg and Byrne, 2019). To get around these problems, constrained decoding techniques have been proposed, most of which artificially impose sparsity on softmax prediction. For example, Fan
et al. (2018) propose to sample from the top- $k$  probable words, and Holtzman et al. (2020) 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 (Martins and Astudillo, 2016) and its generalization  $\alpha$ -entmax (Peters et al., 2019). 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  $\alpha$ -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  $\alpha$ -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  $\alpha$ -entmax in the NMT tasks,
- works as fast as softmax during training and at inference,
- gives the same training dynamics as  $\alpha$ -entmax (in training steps).
The most surprising thing is that such a transformation is simply a shifted ReLU raised to power  $\frac{1}{\alpha - 1}$ , which we call  $\alpha$ -ReLU.
2We also compare against an approximate version which only performs sorting on the highest values of the logits.
<sup>1</sup>The authors called this phenomenon the cat got your tongue problem.
Findings of the Association for Computational Linguistics: NAACL 2022, pages 1142 - 1158 July 10-15, 2022 ©2022 Association for Computational Linguistics
The rest of the paper is organised as follows. In Sect. 2 we motivate the choice of  $\alpha$ -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  $\alpha$ -ReLU in the NMT task. Sect. 4 is devoted to a comparative analysis of  $\alpha$ -ReLU and  $\alpha$ -entmax in terms of sparsity, ability to solve the empty translation problem, and training dynamics.
2  $\alpha$ -ReLU at Output
Our departure point is the  $\alpha$ -entmax transformation of Peters et al. (2019) which can be defined for  $\mathbf{z} \in \mathbb{R}^d$  as
$$
\alpha \text {- e n t m a x} _ {i} (\mathbf {z}) = \left[ (\alpha - 1) z _ {i} - \tau (\mathbf {z}) \right] _ {+} ^ {\frac {1}{\alpha - 1}}, \tag {1}
$$
where  $[x]_{+} := \max \{x, 0\}$ , and  $\tau: \mathbb{R}^d \to \mathbb{R}$  is the (unique) function that satisfies  $\sum_{j}[(\alpha - 1)z_j - \tau(\mathbf{z})]_{+}^{\frac{1}{\alpha - 1}} = 1$  for any  $\mathbf{z}$ . It is this threshold  $\tau$  that makes the computation of  $\alpha$ -entmax slow, because one needs to sort the components of  $\mathbf{z}$  to find  $\tau$  (Peters et al., 2019, Alg. 2).
As we can see, the threshold  $\tau$  is only needed to ensure that  $\alpha$ -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
$$
\alpha \text {- R e L U} _ {i} (\mathbf {z}) := [ (\alpha - 1) z _ {i} - \tau ] _ {+} ^ {\frac {1}{\alpha - 1}}, \tag {2}
$$
where  $\tau$  is a constant that does not depend on  $\mathbf{z}$ . In order to force  $\alpha\text{-ReLU}(\mathbf{z})$  applied to the logits  $\mathbf{z}$  to converge to the one-hot vector  $\mathbf{e}_y$  of the gold label  $y$  we need to adjust the corresponding loss. This can easily be done by feeding the logits  $\mathbf{z}$  and the output  $\alpha\text{-ReLU}(\mathbf{z})$  into the following loss, which we call  $\alpha\text{-ReLU}$  loss.
$$
\begin{array}{l} \ell (\mathbf {z}, y) = \left(\alpha - \operatorname {R e L U} (\mathbf {z}) - \mathbf {e} _ {y}\right) ^ {\top} \left(\mathbf {z} - \frac {\tau}{\alpha - 1} \mathbf {1}\right) \\ + \mathrm {H} _ {\alpha} [ \alpha - \operatorname {R e L U} (\mathbf {z}) ], \tag {3} \\ \end{array}
$$
where  $\mathrm{H}_{\alpha}[\mathbf{p}]\coloneqq \frac{1}{\alpha(\alpha - 1)}\left(1 - \sum_{j}p_{j}^{\alpha}\right),\alpha \neq 1$  is the Tsallis  $\alpha$  entropy (Tsallis, 1988), and  $\mathbf{1}\coloneqq (1,\dots ,1)\in \mathbb{R}^{d}$  is a vector of ones. The rationale for coupling  $\alpha$  -ReLU with the loss (3) is the following
Lemma 1. For any  $\tau \in \mathbb{R}$ , the gradient of the  $\alpha$ -ReLU loss (3) is given by
$$
\nabla_ {\mathbf {z}} \ell (\mathbf {z}, y) = \alpha \mathrm {- R e L U} (\mathbf {z}) - \mathbf {e} _ {y}.
$$
Proof. The proof is in Appendix B.1.
By Lemma 1, gradient-based minimization of  $\ell$  indeed forces  $\alpha\text{-ReLU}(\mathbf{z}) \to \mathbf{e}_y$ . Notice that this is similar to what happens when the softmax normalization is coupled with the cross-entropy loss or when  $\alpha$ -entmax is coupled with the entmax loss. In both cases differentiating the loss with respect to logits gives  $\mathbf{p} - \mathbf{e}_y$ , where  $\mathbf{p}$  is either  $\mathrm{softmax}(\mathbf{z})$  or  $\alpha$ -entmax( $\mathbf{z}$ ) (Martins and Astudillo, 2016; Peters et al., 2019).
Remark. Recall that  $\alpha$ -entmax is a generalization of sparsemax. For example, 2-entmax is essentially sparsemax, and for  $\alpha \in (1,2)$  we get a smoothed version of sparsemax. Similarly,  $\alpha$ -ReLU is a kind of generalization of ReLU. So, the standard ReLU is 2-ReLU (with  $\tau = 0$ ), and for  $\alpha \in (1,2)$  we get a smoothed ReLU (see Fig. 1).
Figure 1: The graph of  $\alpha$ -ReLU(x) for several  $\alpha \in (1,2]$ , with  $\tau = 0$ . 2-ReLU is a standard  $\mathrm{ReLU}(x) := [x]_{+}$ .
In theory, nothing prevents  $\alpha$ -ReLU from learning what  $\alpha$ -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  $\alpha$ -ReLU empirically with  $\alpha$ -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  $\alpha$ -ReLU as drop-in replacement for  $\alpha$ -entmax.
We test  $\alpha$ -ReLU at output in a neural machine translation task (Sutskever et al., 2014), which is essentially a conditional text generation task. Compared to open-ended text generation, there is a
Output Transform	Loss	IWSLT De→En	WMT En→De	WMT En→Ru
softmax	cross-entropy	35.3	28.7	22.4
sparsemax	sparsemax loss	35.5	26.6	19.6
1.5-entmax	1.5-entmax loss	36.6	28.6	23.9
1.5-entmax (k=100)	1.5-entmax loss	36.7	28.4	23.7
1.5-ReLU	1.5-ReLU loss	37.3	28.6	24.6
# Trainable parameters	47M	75M	75M
Table 1: NMT results: comparison of softmax, sparsemax, 1.5-Entmax and the proposed 1.5-ReLU as the output transformations in the Transformer NMT model. Reported is detokenized test BLEU.
clearer metric of the quality of the generated text—the BLEU score (Papineni et al., 2002). As in open-ended text generation, at each prediction step, the NMT system needs to make a choice from all words (subwords) of the vocabulary, the size of which can reach several tens of thousands. Therefore, the sparsity of the output distribution becomes critical in such setups, since it can explicitly prevent the occurrence of most of the words that are inappropriate in the context.
Data. We conduct experiments on three datasets of varied sizes:
- IWSLT'14 De→En (Cettolo et al.), 172K training examples,
- WMT'14 En→De (Bojar et al., 2014), 4.5M training examples,
WMT'13 En→Ru (Bojar et al., 2013), 1.3M training examples. $^{3}$
We preprocess all datasets using the byte pair encoding algorithm (Sennrich et al., 2016) with 10K merge operations on IWSLT, 40K merge operations on WMT En→De, and 60K merge operations on WMT En→Ru. We report detokenized case-sensitive BLEU with SacreBLEU (Post, 2018).<sup>4</sup>
Hyperparameters  $\alpha$  and  $\tau$ . In all experiments we set  $\alpha = 1.5$ , because this value was recommended by Peters et al. (2019); Peters and Martins (2021) as the middle ground between  $\alpha = 1$  (softmax) and  $\alpha = 2$  (sparsemax).
The value for  $\tau$  is chosen as follows: we run the first batch through a non-trained neural network,
which has 1.5-entmax at the output, in the forward direction and determine the average  $\tau$  value across the batch. This value is then used to train the 1.5-ReLU network. Our preliminary experiments have shown that 1.5-ReLU convergence is sensitive to the  $\tau$  value, and that having output close to the probability distribution early in the learning phase works well with the rest of hyperparameters which are set to their default values.
Training. We trained the Transformer Base (Vaswani et al., 2017) using the OpenNMT-py 2.0 toolkit (Klein et al., 2017). Optimization details are in Appendix A.
The results are given in Table 1. Reported are test BLEU scores for best checkpoints which are selected based on validation BLEU. We observe that the 1.5-ReLU performs on par with 1.5-entmax or better, while sparsemax is inferior to all others.
Training Time. Fig. 2&3 show the training dynamics in training steps and in wall time on WMT'14 En→De. Despite the closeness of performance in intermediate steps and at the end of training, we see that on the larger datasets 1.5-entmax is slower in wall time than softmax and 1.5-ReLU.
To speed up the learning process, Peters et al. (2019) recommended limiting the number of sorted logits in the  $\alpha$ -entmax to the  $k$  largest logits. We tried this using  $k = 100$ , which is the default value in the author's implementation of  $\alpha$ -entmax. The resulting training dynamics are shown as dashed curves in Fig. 2&3. As we can see, partial sorting indeed speeds up the learning process, and at the same time does not harm the quality of the translation compared to  $\alpha$ -entmax with full sorting. But in the end, learning is still slower than in the case
<https://github.com/deep-spin/entmax>
3We did not use the Yandex 1M Parallel Corpus because of its license restrictions.  $^4 BLEU +$  case.mixed+lang.en- de+numrefs.1+smooth.exp+tok.13a+version.1.5.1
Figure 2: Training dynamics in training steps.
Figure 3: Training dynamics in absolute time. 1.5-entmax ( $k = 100$ ) is a variant of 1.5-entmax in which sorting is performed only for the largest  $k = 100$  logits.
Figure 4: Normalized inference for WMT En  $\rightarrow$  Ru with different beam sizes.
of 1.5-ReLU. Of course, one can try to select such  $k$  that the speed of calculating the 1.5-entmax will be as close as possible to the speed of 1.5-ReLU without losing quality, but this requires additional efforts on the part of the user, and this must be done for each case separately. Also note that both 1.5-entmaxes (with full and partial sorting) cannot learn the English-Russian data set as well as 1.5-ReLU.
In this regard, 1.5-ReLU does not require additional fine-tuning, converges as fast as softmax in absolute time and performs on par or better. Thus 1.5-ReLU combines all three desired properties: computation speed, task performance, and sparsity of output.
Inference Time. We measured inference time of translating the WMT  $\mathrm{En} \rightarrow \mathrm{Ru}$  test data with the different strategies and with different beam sizes. The results—normalized by the smallest value—are shown in Fig. 4. As can be seen the relative difference seems independent of the beam size: softmax is almost twice faster than 1.5-entmax (with full sorting over the logits). Even though the softmax version is optimized through the softmax CUDA kernel, it performs equivalent to the 1.5-ReLU model in terms of computation speed.
4.1 Empty Translations
We remind the reader that the cat got your tongue problem (Stahlberg and Byrne, 2019) is one of the main motivations for using sparse transformations when generating text. As Peters and Martins (2021) have shown, 1.5-entmax successfully tackles this problem by significantly lowering the proportion of cases where an empty string is more likely than the beam search hypothesis. For 1.5-ReLU, we also calculated this proportion, and compared it with the proportions for softmax and sparsemax (Table 2). As we see, 1.5-ReLU also successfully tackles the cat got your tongue problem.
Figure 5: Sparsity as proportion of zero components after applying 1.5-ReLU and 1.5-entmax, test sets.
Figure 6: Sparsity on training set.
Output Transform	IWSLT De→En	WMT En→De	WMT En→Ru
softmax	7.5%	29.8%	31.7%
sparsemax	0%	0.03%	0%
1.5-entmax	0%	0.2%	0%
1.5-ReLU	0%	0.3%	0.1%
Table 2: Percentage of development set examples for which the model assigns higher probability to the empty string than to the beam-decoded hypothesis.
4.2 Sparsity
To compare the sparsity of 1.5-ReLU and 1.5-entmax we depict in Fig. 5 the distributions of the number of zero components after applying these transformations (recall that for softmax all components are always nonzero). Since we constructed the  $\alpha$ -ReLU in such way that it mimics the  $\alpha$ -entmax (at least in the early stages of training), we expected that these two transformations would have similar properties, including sparsity. However, this is not the case: as we can see, the 1.5-ReLU is significantly less sparse than the 1.5-entmax. It is noteworthy that lower sparsity in this case correlates with a better performance in the translation task (see Table 1).
A possible explanation for the difference in sparsity levels could be that  $\alpha$ -ReLU, in contrast to  $\alpha$ -entmax, behaves significantly differently on the test set than on the training set. However, this is not the case: for example, comparing the sparsity on the IWSLT training set (Fig. 6), we see that the distributions of non-zero components are almost the same as on the test set for 1.5-ReLU and 1.5-entmax.
Note that the sparsity of  $\alpha$ -ReLU and  $\alpha$ -entmax
is approximately the same at the beginning of training due to how we initialize  $\tau$  in 1.5-ReLU (making it as close as possible to 1.5-entmax's  $\tau$  in the untrained model, Sec. 3.1). However, during training,  $\alpha$ -ReLU's  $\tau$  remains fixed, and the model can only adapt the logits themselves so that  $\alpha$ -ReLU( $\mathbf{z}$ ) converges to the corresponding one-hot vector. At the same time, in  $\alpha$ -entmax,  $\tau(\mathbf{z})$  adapts together with logits  $\mathbf{z}$ . We hypothesize that during training, the entmax's  $\tau(\mathbf{z})$  gradually increases which entails greater sparsity by the end of the training. However, the logits themselves also change during training, so the increase in  $\tau$  may not be the cause of greater sparsity. To find out, we track the dynamics of mean logit norm  $\|\mathbf{z}\|$  and mean  $\tau$  during training for both 1.5-entmax and 1.5-ReLU (Fig. 7). As we can see, the logit sizes grow in both cases.
Figure 7: Evolution of the mean  $\tau (\mathbf{z})$  and  $\| \mathbf{z}\|$  during training for 1.5-entmax and 1.5-ReLU models on IWSLT'14 En  $\rightarrow$  De.
At the same time, the 1.5-entmax's  $\tau(\mathbf{z})$  increases following the logit size, while the 1.5-ReLU's  $\tau$  remains constant. From this we conclude that the sparsity of 1.5-entmax is inevitably less than the sparsity of 1.5-ReLU.
4.3 Impact of  $\tau$
The selection of  $\tau$  was described in Section 3.1. However, the question arises: does the described approach lead to the choice of the optimal  $\tau$ ? To find out, we trained the  $\alpha$ -ReLU models for  $\tau \in \{0,0.1,0.2,\dots,0.9,1,2,5,10\}$  on the IWSLT data. Note that all of these  $\tau$ 's have led to almost the same result at the end of the training (as predicted by Lemma 1). In Fig. 8, we present the dynamics of early training only for  $\tau \in \{0,0.1,0.2,0.3,5,10\}$ , since the curves for  $\tau \in \{0.4,\dots,0.9,1,2\}$  practically coincided with the optimal curve corresponding to  $\tau = 0.3$ . Note that our  $\tau$  selection method
Figure 8: Impact of  $\tau$  on training dynamics, IWSLT'14 En→De.
gave a value of 0.33, thus we have no evidence against the adequacy of our method.
4.4 Estimation of  $\tau$  without data
On closer inspection, we noticed that the pre-entmax logits in the untrained Transformer model are distributed according to the normal law, regardless of what data is supplied to the input, Shapiro-Wilk test,  $p$ -value  $> 0.15$ . This allows us, using asymptotic theory, to estimate  $\tau$  as
$$
\hat {\tau} = \sqrt {\frac {d _ {\mathrm {m o d e l}}}{2 \left(d _ {\mathrm {m o d e l}} + d _ {\mathrm {v o c a b}}\right)}} \cdot \Phi^ {- 1} \left(1 - p ^ {*}\right), \tag {4}
$$
where  $d_{\mathrm{model}}$  is the size of hidden representations,  $d_{\mathrm{vocab}}$  is the vocabulary size for a target language,  $\Phi^{-1}(\cdot)$  is the probit function and  $p^*$  is the solution of a non-linear equation that involves functions related to the standard normal distribution (see Appendix B.2 for details). Table 3 compares the  $\tilde{\tau}$  calculated by running data through an untrained model with the estimate  $\hat{\tau}$  obtained from (4). As we can see,  $\hat{\tau}$  practically coincides with  $\tilde{\tau}$  with an
IWSLT'14 De→En	WMT'14 En→De	WMT'13 En→Ru
dmodel	512	512	512
dvocab	10,000	40,000	60,000
p*	.0184	.0171	.0169
\(\tilde{\tau}\)	.33	.17	.14
\(\hat{\tau}\)	.33	.17	.14
Table 3: Estimating threshold of 1.5-entmax:  $\tilde{\tau}$  is a value obtained by running a data through an untrained model;  $\hat{\tau}$  is an estimate based on asymptotic theory, i.e. without running the data through the model.
accuracy of two decimal places. Unfortunately, the formula (4) is not universal: it is only true for the Transformer architecture.
4.5 Self-normalization
The attentive reader may have noticed that the output of  $\alpha$ -ReLU is not normalized, i.e. the components of  $\alpha$ -ReLU( $\mathbf{z}$ ) do not have to sum up to 1. Accordingly, the question arises: how correct is it to compare translation scores at different steps of the beam-search decoding if the conditional probabilities are not normalized? However, the comparison is possible if the  $\alpha$ -ReLU( $\mathbf{z}$ ) components add up to approximately the same number, i.e. if the model is self-normalizing. To check this, we ran the trained  $\alpha$ -ReLU model on the IWSLT and WMT'14 test sets, and looked at the distribution of  $\sum_{i} \alpha$ -ReLU $_i(\mathbf{z})$  at each decoding step. The results are shown in Fig. 9. As we can see, the sum of the
Figure 9: Distribution of the sum of  $\alpha$ -ReLU(z) components across the IWSLT'14 and WMT'14 test sets:  $\alpha$ -ReLU self-normalizes.
$\alpha$ -ReLU  $(\mathbf{z})$  components concentrates well around its mean  $\approx 1.24$  (IWSLT) and 1.09 (WMT'14), which might indicate that the model indeed has a self-normalization property.
4.6 Training Dynamics
As we noted in Sect. 3.2, the training dynamics are similar in all three cases (softmax, 1.5-
entmax, 1.5-ReLU) when time is measured in training steps. Here we attempt to explain this phenomenon through the recently proposed Neural Tangent Kernel (NTK) approach of Jacot et al. (2018). Roughly speaking, the NTK theory suggests that a sufficiently wide neural network trains like a kernel regression. We use this theory to show (in Appendix B.3) that in all three cases the logits  $\mathbf{z}(x,t)$  for a training instance  $x$  at a training step  $t$  evolve (approximately) according to the same differential equation
$$
\frac {d \mathbf {z}}{d t} = - \mathbb {E} _ {\left(x ^ {\prime}, y ^ {\prime}\right)} \left[ \mathbf {K} _ {\sigma} \left(x, x ^ {\prime}\right) \cdot \left(\sigma \left(\mathbf {z} ^ {\prime}\right) - \mathbf {e} _ {y ^ {\prime}}\right) \right], \tag {5}
$$
where expectation is over training examples  $(x^{\prime},y^{\prime})$ ,  $\sigma (\cdot)$  is one of the transformations considered (softmax,  $\alpha$ -entmax, or  $\alpha$ -ReLU), and  $\mathbf{K}_{\sigma}(x,x^{\prime})\in \mathbb{R}^{d\times d}$  is a positive semi-definite matrix that depends on  $\sigma$ . The Equation (5) is a non-linear matrix differential equation which in general cannot be solved analytically. However, it has an equilibrium point  $\mathbf{z}(x,t)$  such that  $\mathbb{E}_{(x',y')}[\mathbf{K}_{\sigma}(x,x')\cdot (\sigma (\mathbf{z}') - \mathbf{e}_{y'})] = \mathbf{0}$ , thus its solution converges to this point as  $t\to \infty$ . This similarity in the evolution of  $\sigma (\mathbf{z})$  implies the similarity in the evolution of the performance metric—such as BLEU—across all three transformations.
4.7 Human Evaluation
Although the BLEU metric (Papineni et al., 2002) has stood the test of time, it is still an automated assessment of translation quality. To double-check the reliability of the results from Table 1, we decided to manually evaluate the translations from the WMT'13 En→Ru test split. To do this, we followed the human evaluation setup from (Berard et al., 2019). We formed two random samples of 135 instances each and gave them to two annotators. 45 instances were shared across two samples in order to calculate inter-annotator agreement. Each instance consists of an original sentence in English and 4 candidate translations into Russian (reference, softmax, entmax,  $\alpha$ -ReLU). The annotators were to rate each translation on a 4-point scale. For annotation instructions, see Appendix C.
The order of candidate translations was shuffled for each instance, so the annotators did not know which sentence is from which model. Nevertheless, the annotator always had a good chance of guessing which translation was the reference one, due to the large difference in quality between human and machine translation.
Model	Avg. Score	Std. Dev.
Reference	3.9	0.30
Softmax	3.3	0.75
1.5-entmax	3.2	0.74
1.5-ReLU	3.3	0.74