Sparse Sequence-to-Sequence Models

Introduction

Attention-based sequence-to-sequence (seq2seq) models have proven useful for a variety of NLP applications, including machine translation Bahdanau et al. (2015); Vaswani et al. (2017), speech recognition Chorowski et al. (2015), abstractive summarization Chopra et al. (2016), and morphological inflection generation Kann and Schütze (2016), among others. In part, their strength comes from their flexibility: many tasks can be formulated as transducing a source sequence into a target sequence of possibly different length.

However, conventional seq2seq models are dense: they compute both attention weights and output probabilities with the softmax function (Bridle, 1990), which always returns positive values. This results in dense attention alignments, in which each source position is attended to at each target position, and in dense output probabilities, in which each vocabulary type always has nonzero probability of being generated. This contrasts with traditional statistical machine translation systems, which are based on sparse, hard alignments, and decode by navigating through a sparse lattice of phrase hypotheses. Can we transfer such notions of sparsity to modern neural architectures? And if so, do they improve performance?

In this paper, we provide an affirmative answer to both questions by proposing neural sparse seq2seq models that replace the softmax transformations (both in the attention and output) by sparse transformations. Our innovations are rooted in the recently proposed sparsemax transformation (Martins and Astudillo, 2016) and Fenchel-Young losses (Blondel et al., 2019). Concretely, we consider a family of transformations (dubbed $\alpha$-entmax), parametrized by a scalar $\alpha$, based on the Tsallis entropies (Tsallis, 1988). This family includes softmax ($\alpha=1$) and sparsemax ($\alpha=2$) as particular cases. Crucially, entmax transforms are sparse for all $\alpha>1$.

Our models are able to produce both sparse attention, a form of inductive bias that increases focus on relevant source words and makes alignments more interpretable, and sparse output probabilities, which together with auto-regressive models can lead to probability distributions that are nonzero only for a finite subset of all possible strings. In certain cases, a short list of plausible outputs can be enumerated without ever exhausting the beam (Figure 1), rendering beam search exact. Sparse output seq2seq models can also be used for adaptive, sparse next word suggestion (Figure 2).

Overall, our contributions are as follows:

We propose an entmax sparse output layer, together with a natural loss function. In large-vocabulary settings, sparse outputs avoid wasting probability mass on unlikely outputs, substantially improving accuracy. For tasks with little output ambiguity, entmax losses, coupled with beam search, can often produce exact finite sets with only one or a few sequences. To our knowledge, this is the first study of sparse output probabilities in seq2seq problems.

We construct entmax sparse attention, improving interpretability at no cost in accuracy. We show that the entmax gradient has a simple form (Proposition 2), revealing an insightful missing link between softmax and sparsemax.

We derive a novel exact algorithm for the case of 1.5-entmax, achieving processing speed close to softmax on the GPU, even with large vocabulary sizes. For arbitrary $\alpha$, we investigate a GPU-friendly approximate algorithm.Our standalone Pytorch entmax implementation is available at https://github.com/deep-spin/entmax.

We experiment on two tasks: one character-level with little ambiguity (morphological inflection generation) and another word-level, with more ambiguity (neural machine translation). The results show clear benefits of our approach, both in terms of accuracy and interpretability.

Background

The underlying architecture we focus on is an RNN-based seq2seq with global attention and input-feeding (Luong et al., 2015). We provide a brief description of this architecture, with an emphasis on the attention mapping and the loss function.

Encoder.

Given an input sequence of tokens $\bm{x}\coloneqq[x_{1},\dots,x_{J}]$, the encoder applies an embedding lookup followed by $K$ layered bidirectional LSTMs (Hochreiter and Schmidhuber, 1997), resulting in encoder states $[\bm{h}_{1},\dots,\bm{h}_{J}]$.

Decoder.

The decoder generates output tokens $y_{1},\ldots,y_{T}$, one at a time, terminated by a stop symbol. At each time step $t$, it computes a probability distribution for the next generated word $y_{t}$, as follows. Given the current state $\bm{s}_{t}$ of the decoder LSTM, an attention mechanism (Bahdanau et al., 2015) computes a focused, fixed-size summary of the encodings $[\bm{h}_{1},\dots,\bm{h}_{J}]$, using $\bm{s}_{t}$ as a query vector. This is done by computing token-level scores $z_{j}\coloneqq\bm{s}_{t}^{\top}\bm{W}^{(z)}\bm{h}_{j}$, then taking a weighted average

The contextual output is the non-linear combination $\bm{o}_{t}\coloneqq\operatorname*{\mathsf{tanh}}(\bm{W}^{(o)}[\bm{s}_{t};\bm{c}_{t}]+\bm{b}^{(o)})$, yielding the predictive distribution of the next word

The output $\bm{o}_{t}$, together with the embedding of the predicted $\widehat{y}_{t}$, feed into the decoder LSTM for the next step, in an auto-regressive manner. The model is trained to maximize the likelihood of the correct target sentences, or equivalently, to minimize

which maps a vector of scores $\bm{z}$ into a probability distribution (i. e., a vector in $\triangle^{d}$). As seen above, the $\operatorname*{\mathsf{softmax}}$ mapping plays two crucial roles in the decoder: first, in computing normalized attention weights (Eq. 1), second, in computing the predictive probability distribution (Eq. 2). Since $\operatorname*{\mathsf{exp}}\gneq 0$, softmax never assigns a probability of zero to any word, so we may never fully rule out non-important input tokens from attention, nor unlikely words from the generation vocabulary. While this may be advantageous for dealing with uncertainty, it may be preferrable to avoid dedicating model resources to irrelevant words. In the next section, we present a strategy for differentiable sparse probability mappings. We show that our approach can be used to learn powerful seq2seq models with sparse outputs and sparse attention mechanisms.

Sparse Attention and Outputs

Since Eq. 5 is a projection onto $\triangle^{d}$, which tends to yield sparse solutions, the predictive distribution $\bm{p}^{\star}\coloneqq\operatorname*{\mathsf{sparsemax}}(\bm{z})$ is likely to assign exactly zero probability to low-scoring choices. They also propose a corresponding loss function to replace the negative log likelihood loss $\mathsf{L}_{\operatorname*{\mathsf{softmax}}}$ (Eq. 3):

This loss is smooth and convex on $\bm{z}$ and has a margin: it is zero if and only if $\bm{z}_{y}\geq\bm{z}_{y^{\prime}}+1$ for any $y^{\prime}\neq y$ (Martins and Astudillo, 2016, Proposition 3). Training models with the sparsemax loss requires its gradient (cf. Appendix A.2):

For using the sparsemax mapping in an attention mechanism, Martins and Astudillo (2016) show that it is differentiable almost everywhere, with

where $s_{j}=1$ if $p^{\star}_{j}>0$, otherwise $s_{j}=0$.

At first glance, sparsemax appears very different from softmax, and a strategy for producing other sparse probability mappings is not obvious. However, the connection becomes clear when considering the variational form of softmax (Wainwright and Jordan, 2008):

where $\mathsf{H}^{\textsc{S}}(\bm{\bm{p}})\coloneqq-\sum_{j}p_{j}\operatorname*{\mathsf{log}}p_{j}$ is the well-known Gibbs-Boltzmann-Shannon entropy with base $e$.

Likewise, letting $\mathsf{H}^{\textsc{G}}(\bm{p})\coloneqq\frac{1}{2}\sum_{j}p_{j}(1-p_{j})$ be the Gini entropy, we can rearrange Eq. 5 as

crystallizing the connection between softmax and sparsemax: they only differ in the choice of entropic regularizer.

2 A new entmax mapping and loss family

The parallel above raises a question: can we find interesting interpolations between softmax and sparsemax? We answer affirmatively, by considering a generalization of the Shannon and Gini entropies proposed by Tsallis (1988): a family of entropies parametrized by a scalar $\alpha>1$ which we call Tsallis $\alpha$-entropies:

This family is continuous, i. e., $\operatorname*{\mathsf{lim}}_{\alpha\rightarrow 1}\mathsf{H}^{\textsc{T}}_{\alpha}(\bm{p})=\mathsf{H}^{\textsc{S}}(\bm{p})$ for any $\bm{p}\in\triangle^{d}$ (cf. Appendix A.1). Moreover, $\mathsf{H}^{\textsc{T}}_{2}\equiv\mathsf{H}^{\textsc{G}}$. Thus, Tsallis entropies interpolate between the Shannon and Gini entropies. Starting from the Tsallis entropies, we construct a probability mapping, which we dub entmax:

and, denoting $\bm{p}^{\star}\coloneqq\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}(\bm{z})$, a loss function

The motivation for this loss function resides in the fact that it is a Fenchel-Young loss (Blondel et al., 2019), as we briefly explain in Appendix A.2. Then, $\mathop{\mathsf{1}\textnormal{-}\mathsf{entmax}}\equiv\operatorname*{\mathsf{softmax}}$ and $\mathop{\mathsf{2}\textnormal{-}\mathsf{entmax}}\equiv\operatorname*{\mathsf{sparsemax}}$. Similarly, $\mathsf{L}_{1}$ is the negative log likelihood, and $\mathsf{L}_{2}$ is the sparsemax loss. For all $\alpha>1$, entmax tends to produce sparse probability distributions, yielding a function family continuously interpolating between softmax and sparsemax, cf. Figure 3. The gradient of the entmax loss is

Tsallis entmax losses have useful properties including convexity, differentiability, and a hinge-like separation margin property: the loss incurred becomes zero when the score of the correct class is separated by the rest by a margin of $\nicefrac{{1}}{{\alpha-1}}$. When separation is achieved, $\bm{p}^{\star}=\bm{e}_{y}$ (Blondel et al., 2019). This allows entmax seq2seq models to be adaptive to the degree of uncertainty present: decoders may make fully confident predictions at “easy” time steps, while preserving sparse uncertainty when a few choices are possible (as exemplified in Figure 2).

Tsallis entmax probability mappings have not, to our knowledge, been used in attention mechanisms. They inherit the desirable sparsity of sparsemax, while exhibiting smoother, differentiable curvature, whereas sparsemax is piecewise linear.

3 Computing the entmax mapping

Whether we want to use $\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}$ as an attention mapping, or $\mathsf{L}_{\alpha}$ as a loss function, we must be able to efficiently compute $\bm{p}^{\star}=\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}(\bm{z})$, i. e., to solve the maximization in Eq. 10. For $\alpha=1$, the closed-form solution is given by Eq. 4. For $\alpha>1$, given $\bm{z}$, we show that there is a unique threshold $\tau$ such that (Appendix C.1, Lemma 2):

i. e., entries with score $z_{j}\leq\nicefrac{{\tau}}{{\alpha-1}}$ get zero probability. For sparsemax ($\alpha=2$), the problem amounts to Euclidean projection onto $\triangle^{d}$, for which two types of algorithms are well studied:

exact, based on sorting Held et al. (1974); Michelot (1986),

iterative, bisection-based Liu and Ye (2009).

The bisection approach searches for the optimal threshold $\tau$ numerically. Blondel et al. (2019) generalize this approach in a way applicable to $\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}$. The resulting algorithm is (cf. Appendix C.1 for details):

Algorithm 1 works by iteratively narrowing the interval containing the exact solution by exactly half. Line 7 ensures that approximate solutions are valid probability distributions, i. e., that $\bm{p}^{\star}\in\triangle^{d}$.

Although bisection is simple and effective, an exact sorting-based algorithm, like for sparsemax, has the potential to be faster and more accurate. Moreover, as pointed out by Condat (2016), when exact solutions are required, it is possible to construct inputs $\bm{z}$ for which bisection requires arbitrarily many iterations. To address these issues, we propose a novel, exact algorithm for 1.5-entmax, halfway between softmax and sparsemax.

We give a full derivation in Appendix C.2. As written, Algorithm 2 is $O(d\operatorname*{\mathsf{log}}d)$ because of the sort; however, in practice, when the solution $\bm{p}^{\star}$ has no more than $k$ nonzeros, we do not need to fully sort $\bm{z}$, just to find the $k$ largest values. Our experiments in §4.2 reveal that a partial sorting approach can be very efficient and competitive with softmax on the GPU, even for large $d$. Further speed-ups might be available following the strategy of Condat (2016), but our simple incremental method is very easy to implement on the GPU using primitives available in popular libraries (Paszke et al., 2017).

Our algorithm resembles the aforementioned sorting-based algorithm for projecting onto the simplex (Michelot, 1986). Both algorithms rely on the optimality conditions implying an analytically-solvable equation in $\tau$: for sparsemax ($\alpha=2$), this equation is linear, for $\alpha=1.5$ it is quadratic (Eq. 40 in Appendix C.2). Thus, exact algorithms may not be available for general values of $\alpha$.

4 Gradient of the entmax mapping

The following result shows how to compute the backward pass through $\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}$, a requirement when using $\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}$ as an attention mechanism.

Let $\alpha\geq 1$. Assume we have computed $\bm{p}^{\star}=\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}(\bm{z})$, and define the vector

Then, $\displaystyle\frac{\partial\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}(\bm{z})}{\partial\bm{z}}=\operatorname*{\mathsf{diag}}(\bm{s})-\frac{1}{\|\bm{s}\|_{1}}~{}\bm{ss}^{\top}.$

Proof: The result follows directly from the more general Proposition 2, which we state and prove in Appendix B, noting that $\left(\frac{t^{\alpha}-t}{\alpha(\alpha-1)}\right)^{\prime\prime}=t^{\alpha-2}$.

The gradient expression recovers the softmax and sparsemax Jacobians with $\alpha=1$ and $\alpha=2$, respectively (Martins and Astudillo, 2016, Eqs. 8 and 12), thereby providing another relationship between the two mappings. Perhaps more interestingly, Proposition 1 shows why the sparsemax Jacobian depends only on the support and not on the actual values of $\bm{p}^{\star}$: the sparsemax Jacobian is equal for $\bm{p}^{\star}=[.99,.01,0]$ and $\bm{p}^{\star}=[.5,.5,0]$. This is not the case for $\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}$ with $\alpha\neq 2$, suggesting that the gradients obtained with other values of $\alpha$ may be more informative. Finally, we point out that the gradient of entmax losses involves the entmax mapping (Eq. 12), and therefore Proposition 1 also gives the Hessian of the entmax loss.

Experiments

The previous section establishes the computational building blocks required to train models with entmax sparse attention and loss functions. We now put them to use for two important NLP tasks, morphological inflection and machine translation. These two tasks highlight the characteristics of our innovations in different ways. Morphological inflection is a character-level task with mostly monotonic alignments, but the evaluation demands exactness: the predicted sequence must match the gold standard. On the other hand, machine translation uses a word-level vocabulary orders of magnitude larger and forces a sparse output layer to confront more ambiguity: any sentence has several valid translations and it is not clear beforehand that entmax will manage this well.

Despite the differences between the tasks, we keep the architecture and training procedure as similar as possible. We use two layers for encoder and decoder LSTMs and apply dropout with probability 0.3. We train with Adam (Kingma and Ba, 2015), with a base learning rate of 0.001, halved whenever the loss increases on the validation set. We use a batch size of 64. At test time, we select the model with the best validation accuracy and decode with a beam size of 5. We implemented all models with OpenNMT-py (Klein et al., 2017).Our experiment code is at https://github.com/deep-spin/OpenNMT-entmax.

In our primary experiments, we use three $\alpha$ values for the attention and loss functions: $\alpha=1$ (softmax), $\alpha=1.5$ (to which our novel Algorithm 2 applies), and $\alpha=2$ (sparsemax). We also investigate the effect of tuning $\alpha$ with increased granularity.

The goal of morphological inflection is to produce an inflected word form (such as “drawn”) given a lemma (“draw”) and a set of morphological tags ({verb, past, participle}). We use the data from task 1 of the CoNLL–SIGMORPHON 2018 shared task (Cotterell et al., 2018). shared task

We train models under two data settings: high (approximately 10,000 samples per language in 86 languages) and medium (approximately 1,000 training samples per language in 102 languages). We depart from previous work by using multilingual training: each model is trained on the data from all languages in its data setting. This allows parameters to be shared between languages, eliminates the need to train language-specific models, and may provide benefits similar to other forms of data augmentation (Bergmanis et al., 2017). Each sample is presented as a pair: the source contains the lemma concatenated to the morphological tags and a special language identification token (Johnson et al., 2017; Peters et al., 2017), and the target contains the inflected form. As an example, the source sequence for Figure 1 is english ␣ verb ␣ participle ␣ past ␣ d ␣ r ␣ a ␣ w. Although the set of inflectional tags is not sequential, treating it as such is simple to implement and works well in practice (Kann and Schütze, 2016). All models use embedding and hidden state sizes of 300. We validate at the end of every epoch in the high setting and only once every ten epochs in medium because of its smaller size.

Accuracy.

Results are shown in Table 1. We report the official metric of the shared task, word accuracy averaged across languages. In addition to the average results of three individual model runs, we use an ensemble of those models, where we decode by averaging the raw probabilities at each time step. Our best sparse loss models beat the softmax baseline by nearly a full percentage point with ensembling, and up to two and a half points in the medium setting without ensembling. The choice of attention has a smaller impact. In both data settings, our best model on the validation set outperforms all submissions from the 2018 shared task except for UZH (Makarov and Clematide, 2018), which uses a more involved imitation learning approach and larger ensembles. In contrast, our only departure from standard seq2seq training is the drop-in replacement of softmax by entmax.

Sparsity.

Besides their accuracy, we observed that entmax models made very sparse predictions: the best configuration in Table 1 concentrates all probability mass into a single predicted sequence in 81% validation samples in the high data setting, and 66% in the more difficult medium setting. When the model does give probability mass to more than one sequence, the predictions reflect reasonable ambiguity, as shown in Figure 1. Besides enhancing interpretability, sparsity in the output also has attractive properties for beam search decoding: when the beam covers all nonzero-probability hypotheses, we have a certificate of globally optimal decoding, rendering beam search exact. This is the case on 87% of validation set sequences in the high setting, and 79% in medium. To our knowledge, this is the first instance of a neural seq2seq model that can offer optimality guarantees.

2 Machine Translation

We now turn to a highly different seq2seq regime in which the vocabulary size is much larger, there is a great deal of ambiguity, and sequences can generally be translated in several ways. We train models for three language pairs in both directions:

IWSLT 2017 German $\leftrightarrow$ English (de$\leftrightarrow$en, Cettolo et al., 2017): training size 206,112.

KFTT Japanese $\leftrightarrow$ English (ja$\leftrightarrow$en, Neubig, 2011): training size of 329,882.

WMT 2016 Romanian $\leftrightarrow$ English (ro$\leftrightarrow$en, Bojar et al., 2016): training size 612,422, diacritics removed (following Sennrich et al., 2016b).

We use byte pair encoding (BPE; Sennrich et al., 2016a) to ensure an open vocabulary. We use separate segmentations with 25k merge operations per language for ro$\leftrightarrow$en and a joint segmentation with 32k merges for the other language pairs. de$\leftrightarrow$en is validated once every 5k steps because of its smaller size, while the other sets are validated once every 10k steps. We set the maximum number of training steps at 120k for ro$\leftrightarrow$en and 100k for other language pairs. We use 500 dimensions for word vectors and hidden states.

Evaluation.

Table 2 shows BLEU scores (Papineni et al., 2002) for the three models with $\alpha\in\{1,1.5,2\}$, using the same value of $\alpha$ for the attention mechanism and loss function. We observe that the 1.5-entmax configuration consistently performs best across all six choices of language pair and direction. These results support the notion that the optimal function is somewhere between softmax and sparsemax, which motivates a more fine-grained search for $\alpha$; we explore this next.

Fine-grained impact of 𝜶𝜶\alpha.

Algorithm 1 allows us to further investigate the marginal effect of varying the attention $\alpha$ and the loss $\alpha$, while keeping the other fixed. We report de$\shortrightarrow$en validation accuracy on a fine-grained $\alpha$ grid in Figure 4. On this dataset, moving from softmax toward sparser attention (left) has a very small positive effect on accuracy, suggesting that the benefit in interpretability does not hurt accuracy. The impact of the loss function $\alpha$ (right) is much more visible: there is a distinct optimal value around $\alpha=1.33$, with performance decreasing for too large values. Interpolating between softmax and sparsemax thus inherits the benefits of both, and our novel Algorithm 2 for $\alpha=1.5$ is confirmed to strike a good middle ground. This experiment also confirms that bisection is effective in practice, despite being inexact. Extrapolating beyond the sparsemax loss ($\alpha>2$) does not seem to perform well.

Sparsity.

In order to form a clearer idea of how sparse entmax becomes, we measure the average number of nonzero indices on the de$\shortrightarrow$en validation set and show it in Table 3. As expected, 1.5-entmax is less sparse than sparsemax as both an attention mechanism and output layer. In the attention mechanism, 1.5-entmax’s increased support size does not come at the cost of much interpretability, as Figure 5 demonstrates. In the output layer, 1.5-entmax assigns positive probability to only 16.13 target types out of a vocabulary of 17,993, meaning that the supported set of words often has an intuitive interpretation. Figure 2 shows the sparsity of the 1.5-entmax output layer in practice: the support becomes completely concentrated when generating a phrase like “the tree of life”, but grows when presenting a list of synonyms (“view”, “look”, “glimpse”, and so on). This has potential practical applications as a predictive translation system (Green et al., 2014), where the model’s support set serves as a list of candidate auto-completions at each time step.

Training time.

Importantly, the benefits of sparsity do not come at a high computational cost. Our proposed Algorithm 2 for 1.5-entmax runs on the GPU at near-softmax speeds (Figure 6). For other $\alpha$ values, bisection (Algorithm 1) is slightly more costly, but practical even for large vocabulary sizes. On de$\shortrightarrow$en, bisection is capable of processing about 10,500 target words per second on a single Nvidia GeForce GTX 1080 GPU, compared to 13,000 words per second for 1.5-entmax with Algorithm 2 and 14,500 words per second with softmax. On the smaller-vocabulary morphology datasets, Algorithm 2 is nearly as fast as softmax.

Related Work

Sparsity in the attention and in the output have different, but related, motivations. Sparse attention can be justified as a form of inductive bias, since for tasks such as machine translation one expects only a few source words to be relevant for each translated word. Dense attention probabilities are particularly harmful for long sequences, as shown by Luong et al. (2015), who propose “local attention” to mitigate this problem. Combining sparse attention with fertility constraints has been recently proposed by Malaviya et al. (2018). Hard attention (Xu et al., 2015; Aharoni and Goldberg, 2017; Wu et al., 2018) selects exactly one source token. Its discrete, non-differentiable nature requires imitation learning or Monte Carlo policy gradient approximations, which drastically complicate training. In contrast, entmax is a differentiable, easy to use, drop-in softmax replacement. A recent study by Jain and Wallace (2019) tackles the limitations of attention probabilities to provide interpretability. They only study dense attention in classification tasks, where attention is less crucial for the final predictions. In their conclusions, the authors defer to future work exploring sparse attention mechanisms and seq2seq models. We believe our paper can foster interesting investigation in this area.

Losses for seq2seq models.

Mostly motivated by the challenges of large vocabulary sizes in seq2seq, an important research direction tackles replacing the cross-entropy loss with other losses or approximations (Bengio and Senécal, 2008; Morin and Bengio, 2005; Kumar and Tsvetkov, 2019). While differently motivated, some of the above strategies (e. g., hierarchical prediction) could be combined with our proposed sparse losses. Niculae et al. (2018) use sparsity to predict interpretable sets of structures. Since auto-regressive seq2seq makes no factorization assumptions, their strategy cannot be applied without approximations, such as in Edunov et al. (2018).

Conclusion and Future Work

We proposed sparse sequence-to-sequence models and provided fast algorithms to compute their attention and output transformations. Our approach yielded consistent improvements over dense models on morphological inflection and machine translation, while inducing interpretability in both attention and output distributions. Sparse output layers also provide exactness when the number of possible hypotheses does not exhaust beam search.

Given the ubiquity of softmax in NLP, entmax has many potential applications. A natural next step is to apply entmax to self-attention (Vaswani et al., 2017). In a different vein, the strong morphological inflection results point to usefulness in other tasks where probability is concentrated in a small number of hypotheses, such as speech recognition.

Acknowledgments

This work was supported by the European Research Council (ERC StG DeepSPIN 758969), and by the Fundação para a Ciência e Tecnologia through contracts UID/EEA/50008/2019 and CMUPERI/TIC/0046/2014 (GoLocal). We thank Mathieu Blondel, Nikolay Bogoychev, Gonçalo Correia, Erick Fonseca, Pedro Martins, Tsvetomila Mihaylova, Miguel Rios, Marcos Treviso, and the anonymous reviewers, for helpful discussion and feedback.

References

Appendix A Background

Recall the definition of the Tsallis family of entropies in Eq. 9 for $\alpha\geq 1$,

This family is continuous in $\alpha$, i. e., $\operatorname*{\mathsf{lim}}_{\alpha\rightarrow 1}\mathsf{H}^{\textsc{T}}_{\alpha}(\bm{p})=\mathsf{H}^{\textsc{T}}_{1}(\bm{p})$ for any $\bm{p}\in\triangle^{d}$. Proof: For simplicity, we rewrite $\mathsf{H}^{\textsc{T}}_{\alpha}$ in separable form:

It suffices to show that $\operatorname*{\mathsf{lim}}_{\alpha\rightarrow 1}h_{\alpha}(t)=h_{1}(t)$ for $t\in$. Let $f(\alpha)\coloneqq t-t^{\alpha}$, and $g(\alpha)\coloneqq\alpha(\alpha-1)$. Observe that $\frac{f(1)}{g(1)}=\nicefrac{{0}}{{0}}$, so we are in an indeterminate case. We take the derivatives of $f$ and $g$:

Note also that, as $\alpha\rightarrow\infty$, the denominator grows unbounded, so $\mathsf{H}^{\textsc{T}}_{\infty}\equiv 0$.

A.2 Fenchel-Young losses

In this section, we recall the definitions and properties essential for our construction of $\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}$. The concepts below were formalized by Blondel et al. (2019) in more generality; we present below a less general version, sufficient for our needs.

This justifies our choice of entmax mapping and loss (Eqs. 10–11), as $\bm{\pi}_{-\mathsf{H}^{\textsc{T}}_{\alpha}}=\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}$ and $\mathsf{L}_{-\mathsf{H}^{\textsc{T}}_{\alpha}}=\mathsf{L}_{\alpha}$.

Zero loss. $L_{\Omega}(\bm{z};\bm{y})=0$ if and only if $\bm{y}=\bm{\pi}_{\Omega}(\bm{z})$, i. e., the prediction is exactly correct.

Convexity. $\mathsf{L}_{\Omega}$ is convex in $\bm{z}$.

Differentiability. $\mathsf{L}_{\Omega}$ is differentiable with $\nabla\mathsf{L}_{\Omega}(\bm{z};\bm{y})=\bm{p}^{\star}-\bm{y}$.

Smoothness. If $\Omega$ is strongly convex, then $L_{\Omega}$ is smooth.

Temperature scaling. For any constant $t>0$, $\mathsf{L}_{t\Omega}(\bm{z};\bm{y})=t\mathsf{L}_{\Omega}(\nicefrac{{\bm{z}}}{{t}};\bm{y})$.

To shed light on the generic probability mapping in Eq. 16, we derive below the optimality conditions characterizing its solution. The optimality conditions are essential not only for constructing algorithms for computing $\bm{p}^{\star}$ (Appendix C), but also for deriving the Jacobian of the mapping (Appendix B). The Lagrangian of the maximization in Eq. 16 is

The subgradient KKT conditions are therefore:

Connection to softmax and sparsemax.

We may now directly see that, when $\Omega(\bm{p})\coloneqq\sum_{j}p_{j}\operatorname*{\mathsf{log}}p_{j}$, Eq. 20 becomes $\operatorname*{\mathsf{log}}p_{j}=\bm{z}_{j}+\nu_{j}-\tau-1$, which can only be satisfied if $p_{j}>0$, thus $\bm{\nu}=\bm{0}$. Then, $p_{j}=\nicefrac{{\operatorname*{\mathsf{exp}}(z_{j})}}{{Z}}$, where $Z\coloneqq\operatorname*{\mathsf{exp}}(\tau+1)$. From Eq. 22, $Z$ must be such that $p_{j}$ sums to 1, yielding the well-known softmax expression. In the case of sparsemax, note that for any $\bm{p}\in\triangle^{d}$, we have

Thus, \displaystyle\operatorname*{\mathsf{argmax}}_{\bm{p}\in\triangle^{d}}\bm{p}^{\top}\bm{z}+\mathsf{H}^{\textsc{G}}(\bm{p})=\operatorname*{\mathsf{argmin}}_{\bm{p}\in\triangle^{d}}{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}0.5\Big{(}}\|\bm{p}\|^{2}-2\bm{p}^{\top}\bm{z}{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\left(+\|\bm{z}\|^{2}\right)\Big{)}}=\operatorname*{\mathsf{argmin}}_{\bm{p}\in\triangle^{d}}\|\bm{p}-\bm{z}\|^{2}.

Appendix B Backward pass for generalized sparse attention mappings

When a mapping $\bm{\pi}_{\Omega}$ is used inside the computation graph of a neural network, the Jacobian of the mapping has the important role of showing how to propagate error information, necessary when training with gradient methods. In this section, we derive a new, simple expression for the Jacobian of generalized sparse $\bm{\pi}_{\Omega}$. We apply this result to obtain a simple form for the Jacobian of $\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}$ mappings.

The proof is in two steps. First, we prove a lemma that shows that Jacobians are zero outside of the support of the solution. Then, completing the result, we characterize the Jacobian at the nonzero indices.

Proof: Since $\Omega$ is strictly convex, the $\operatorname*{\mathsf{argmax}}$ in Eq. 16 is unique. Using Danskin’s theorem (Danskin, 1966), we may write

Since $\Omega$ is strongly convex, the gradient of its conjugate $\Omega^{*}$ is differentiable almost everywhere (Rockafellar, 1970). Moreover, $\frac{\partial{\bm{\pi}_{\Omega}}}{\partial{\bm{z}}}$ is the Hessian of $\Omega^{*}$, therefore it is symmetric, proving the first two claims. Recall the definition of a partial derivative,

Denote by $\bm{p}^{\star}\coloneqq\bm{\pi}_{\Omega}(\bm{z})$. We will show that for any $j$ such that $p^{\star}_{j}=0$, and any $\varepsilon\geq 0$,

It follows that $\operatorname*{\mathsf{lim}}_{\varepsilon\rightarrow 0_{-}}\frac{1}{\varepsilon}\left(\bm{\pi}_{\Omega}(\bm{z}+\varepsilon\bm{e}_{j})_{i}-\bm{\pi}_{\Omega}(\bm{z})_{i}\right)=0.$ If $\bm{\pi}_{\Omega}$ is differentiable at $\bm{z}$, this one-sided limit must agree with the derivative. Otherwise, the sparse one-sided limit is a generalized Jacobian.

Let $\bm{p}^{\star}\coloneqq\bm{\pi}_{\Omega}(\bm{z})$, with strongly convex and differentiable $\Omega$. Denote the support of $\bm{p}^{\star}$ by \mathcal{S}=\big{\{}j\in\{1,\dots,d\}:p_{j}>0\big{\}}. If the second derivative $h_{ij}=\frac{\partial^{2}\Omega}{\partial p_{i}\partial p_{j}}(\bm{p}^{\star})$ exists for any $i,j\in\mathcal{S}$, then

In particular, if $\Omega(\bm{p})=\sum_{j}g(p_{j})$ with $g$ twice differentiable on $(0,1]$, we have

Proof: Lemma 1 verifies that $\frac{\partial(\bm{\pi}_{\Omega})_{i}}{\partial z_{j}}=0$ for $i,j\notin\mathcal{S}$. It remains to find the derivatives w.r.t. $i,j\in\mathcal{S}$. Denote by $\bar{\bm{p}}^{\star},\bar{\bm{z}}$ the restriction of the corresponding vectors to the indices in the support $\mathcal{S}$. The optimality conditions on the support are

where \bm{g}(\bar{\bm{p}})\coloneqq\big{(}\nabla\Omega(\bm{p})\big{)}\big{|}_{\mathcal{S}}, so $\frac{\partial\bm{g}}{\partial\bar{\bm{p}}}(\bar{\bm{p}}^{\star})=\bm{H}$. Differentiating w.r.t. $\bar{\bm{z}}$ at $\bm{p}^{\star}$ yields

Since $\Omega$ is strictly convex, $\bm{H}$ is invertible. From block Gaussian elimination (i. e., the Schur complement),

which can then be used to solve for $\frac{\partial\bar{\bm{p}}}{\partial\bar{\bm{z}}}$ giving

yielding the desired result. When $\Omega$ is separable, $\bm{H}$ is diagonal, with $H_{ii}=g^{\prime\prime}(p^{\star}_{i})$, yielding the simplified expression which completes the proof.

Our result is similar, but simpler than Niculae and Blondel (2017, Proposition 1), especially in the case of separable $\Omega$. Crucially, our result does not require that the second derivative exist outside of the support. As such, unlike the cited work, our result is applicable in the case of $\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}$, where either $g^{\prime\prime}(t)=t^{\alpha-2}$ or its reciprocal may not exist at $t=0$.

Appendix C Algorithms for entmax

The following lemma provides a simplified form for the solution of $\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}$.

For any $\bm{z}$, there exists a unique $\tau^{\star}$ such that

Proof: We use the regularized prediction functions defined in Appendix A.2. From both definitions,

We first note that for all $\bm{p}\in\triangle^{d}$,

From the constant invariance and scaling properties of $\bm{\pi}_{\Omega}$ (Blondel et al., 2019, Proposition 1, items 4–5),

Using (Blondel et al., 2019, Proposition 5), noting that $g^{\prime}(t)=t^{\alpha-1}$ and $(g^{\prime})^{-1}(u)=u^{\nicefrac{{1}}{{\alpha-1}}}$, yields

Uniqueness of $\tau^{\star}$ follows from the fact that $\mathop{\mathsf{\alpha}\textnormal{-}\mathsf{entmax}}$ has a unique solution $\bm{p}^{\star}$, and Eq. 32 implies a one-to-one mapping between $\bm{p}^{\star}$ and $\tau^{\star}$, as long as $\bm{p}^{\star}\in\triangle$.

For $\alpha=1.5$, Lemma 2 implies existence of a unique $\tau^{\star}$ such that

C.2 An exact algorithm for entmax with 𝜶=1.5𝜶1.5\alpha=1.5: Derivation of Algorithm 2.

In this section, we derive an exact, sorting-based algorithm for $\mathop{\mathsf{1.5}\textnormal{-}\mathsf{entmax}}$. The key observation is that the solution can be characterized by the size of its support, $\rho^{\star}=\|\bm{p}^{\star}\|_{0}$. Then, we can simply enumerate all possible values of $\rho\in\{1,\dots,d\}$ until the solution verifies all optimality conditions. The challenge, however, is expressing the threshold $\tau$ as a function of the support size $\rho$; for this, we rely on $\alpha=1.5$.

Exact computation of $\mathop{\mathsf{1.5}\textnormal{-}\mathsf{entmax}}(\bm{z})$

Let $z_{[d]}\leq\dots\leq z_{}$ denote the sorted coordinates of $\bm{z}$, and, for convenience, let $z_{[d+1]}\coloneqq-\infty$. Define the top-$\rho$ mean, unnormalized variance, and induced threshold for $\rho\in\{1,\dots,d\}$ as

for any $\rho$ satisfying $\tau_{\bm{z}}(\rho)\in[z_{[\rho+1]},z_{[\rho]}]$.

Proposition 3 implies the correctness of Algorithm 2. To prove it, we first need the following lemma.

Define $\tau(\rho)$ as in Proposition 3. Then, $\tau$ is non-decreasing, and there exists $\rho_{\textsf{max}}\in\{1,\dots,d\}$ such that $\tau$ is finite for $1\leq\rho\leq\rho_{\textsf{max}}$, and infinite for $\rho>\rho_{\textsf{max}}$.

The proof is slightly more technical, and we defer it to after the proof of the proposition.

First, using Corollary 2.1 we reduce the problem of computing $\mathop{\mathsf{1.5}\textnormal{-}\mathsf{entmax}}$ to

Denote by $\tau^{\star}$ the optimal threshold as defined in the corollary. We will show that $\tau^{\star}=\tau(\rho)$ for any $\rho$ satisfying $\tau(\rho)\in[z_{[\rho+1]},z_{[\rho]}]$, where we assume, for convenience, $z_{[d+1]}=-\infty$. The generic stationarity condition in Eq. 20, applied to the problem in Eq. 34, takes the form

Since $\Omega$ is symmetric, $\bm{\pi}_{\Omega}$ is permutation-preserving (Blondel et al., 2019, Proposition 1, item 1), so we may assume w.l.o.g. that $\bm{z}$ is sorted non-increasingly, i.e., $z_{1}\geq\dots\geq z_{d}$; in other words, $z_{j}=z_{[j]}$. Therefore, the optimal $\bm{p}$ is also non-increasing. Denote by $\rho$ an index such as $p_{j}\geq 0$ for $1\leq j\leq\rho$, and $p_{j}=0$ for $j>\rho$. From the complementary slackness condition (21), $\nu_{j}=0$ for $1\leq j\leq\rho$, thus we may split the stationarity conditions (35) into

For (36) to have solutions, the r.h.s. must be non-negative, i.e., $\tau\leq z_{j}$ for $j\leq\rho$, so $\tau\leq z_{\rho}$. At the same time, from dual feasability (23) we have $\nu_{j}=\tau-z_{j}\geq 0$ for $j>\rho$, therefore

Given $\rho$, we can solve for $\tau$ using (36) and primal feasability (22)

Expanding the squares and dividing by $2\rho$ yields the quadratic equation

Therefore, $\tau^{\star}=\tau(\rho)=M(\rho)-\sqrt{\frac{1-S(\rho)}{\rho}}$ for some $\rho$ verifying (38). It remains to show that any such $\rho$ leads to the same value of $\tau(\rho)$. Pick any $\rho_{1}<\rho_{2}$, both verifying (38). Therefore, $\rho_{1}+1\leq\rho_{2}$ and

thus $\tau(\rho_{1})=\tau(\rho_{2})$, and so any $\rho$ verifying (38) satisfies $\tau^{\star}=\tau(\rho)$, concluding the proof. ∎

Proof of Lemma 3.

Domain of $\tau$. The threshold $\tau(\rho)$ is only finite for \rho\in T\coloneqq\big{\{}\rho\in\{1,\dots,d\}\colon S(\rho)\leq 1\big{\}}, i.e., where $\nicefrac{{(1-S(\rho))}}{{\rho}}\geq 0$. We show there exists $\rho_{\textsf{max}}$ such that $T=\{1,\dots,\rho_{\textsf{max}}\}$. Choose $\rho_{\textsf{max}}$ as the largest index satisfying $S(\rho_{\textsf{max}})\leq 1$. By definition, $\rho>\rho_{\textsf{max}}$ implies $\rho\notin T$. Remark that $S(1)=0$, and $S(\rho+1)-S(\rho)=(\cdot)^{2}\geq 0$. Therefore, $S$ is nondecreasing and, for any $1\leq\rho\leq\rho_{\textsf{max}},0\leq S(\rho)\leq 1$.

We seek the lower bound $\widetilde{\tau}(x^{\star})$ and show that $\widetilde{\tau}(x^{\star})\geq\tau_{\bm{z}}(\rho)$. From (43), this implies $\tau_{\bm{z}}(\rho+1)\geq\tau_{\bm{z}}(\rho)$ and, by transitivity, the monotonicity of $\tau_{\bm{z}}$.

It is easy to verify that the following incremental update expressions hold.

Ignoring the constraint for a moment and setting the gradient to yields the solution

implying $x^{\star}<0$. Squaring both sides and rearranging yields the solution of the unconstrained problem,

We verify that $x^{\star}$ readily satisfies the constraints, thus it is a solution to the minimization in Eq. 45:

Plugging $x^{\star}$ into the objective yields

Therefore, $\widetilde{\tau}(x)\geq\tau_{\bm{z}}(\rho)$ for any valid $x$, proving that $\tau_{\bm{z}}(\rho)\leq\tau_{\bm{z}}(\rho+1)$.