Adaptively Sparse Transformers

Abstract: Attention mechanisms have become ubiquitous in NLP. Recent architectures, notably the Transformer, learn powerful context-aware word representations through layered, multi-headed attention. The multiple heads learn diverse types of word relationships. However, with standard softmax attention, all attention heads are dense, assigning a non-zero weight to all context words. In this work, we introduce the adaptively sparse Transformer, wherein attention heads have flexible, context-dependent sparsity patterns. This sparsity is accomplished by replacing softmax with $\alpha$-entmax: a differentiable generalization of softmax that allows low-scoring words to receive precisely zero weight. Moreover, we derive a method to automatically learn the $\alpha$ parameter -- which controls the shape and sparsity of $\alpha$-entmax -- allowing attention heads to choose between focused or spread-out behavior. Our adaptively sparse Transformer improves interpretability and head diversity when compared to softmax Transformers on machine translation datasets. Findings of the quantitative and qualitative analysis of our approach include that heads in different layers learn different sparsity preferences and tend to be more diverse in their attention distributions than softmax Transformers. Furthermore, at no cost in accuracy, sparsity in attention heads helps to uncover different head specializations.

The Transformer architecture BIBREF0 for deep neural networks has quickly risen to prominence in NLP through its efficiency and performance, leading to improvements in the state of the art of Neural Machine Translation BIBREF1, BIBREF2, as well as inspiring other powerful general-purpose models like BERT BIBREF3 and GPT-2 BIBREF4. At the heart of the Transformer lie multi-head attention mechanisms: each word is represented by multiple different weighted averages of its relevant context. As suggested by recent works on interpreting attention head roles, separate attention heads may learn to look for various relationships between tokens BIBREF5, BIBREF6, BIBREF7, BIBREF8, BIBREF9.
The attention distribution of each head is predicted typically using the softmax normalizing transform. As a result, all context words have non-zero attention weight. Recent work on single attention architectures suggest that using sparse normalizing transforms in attention mechanisms such as sparsemax – which can yield exactly zero probabilities for irrelevant words – may improve performance and interpretability BIBREF12, BIBREF13, BIBREF14. Qualitative analysis of attention heads BIBREF0 suggests that, depending on what phenomena they capture, heads tend to favor flatter or more peaked distributions.
Recent works have proposed sparse Transformers BIBREF10 and adaptive span Transformers BIBREF11. However, the “sparsity" of those models only limits the attention to a contiguous span of past tokens, while in this work we propose a highly adaptive Transformer model that is capable of attending to a sparse set of words that are not necessarily contiguous. Figure FIGREF1 shows the relationship of these methods with ours.
Our contributions are the following:
We introduce sparse attention into the Transformer architecture, showing that it eases interpretability and leads to slight accuracy gains.
We propose an adaptive version of sparse attention, where the shape of each attention head is learnable and can vary continuously and dynamically between the dense limit case of softmax and the sparse, piecewise-linear sparsemax case.
We make an extensive analysis of the added interpretability of these models, identifying both crisper examples of attention head behavior observed in previous work, as well as novel behaviors unraveled thanks to the sparsity and adaptivity of our proposed model.
In NMT, the Transformer BIBREF0 is a sequence-to-sequence (seq2seq) model which maps an input sequence to an output sequence through hierarchical multi-head attention mechanisms, yielding a dynamic, context-dependent strategy for propagating information within and across sentences. It contrasts with previous seq2seq models, which usually rely either on costly gated recurrent operations BIBREF15, BIBREF16 or static convolutions BIBREF17.
Given $n$ query contexts and $m$ sequence items under consideration, attention mechanisms compute, for each query, a weighted representation of the items. The particular attention mechanism used in BIBREF0 is called scaled dot-product attention, and it is computed in the following way:
where $\mathbf {Q} \in \mathbb {R}^{n \times d}$ contains representations of the queries, $\mathbf {K}, \mathbf {V} \in \mathbb {R}^{m \times d}$ are the keys and values of the items attended over, and $d$ is the dimensionality of these representations. The $\mathbf {\pi }$ mapping normalizes row-wise using softmax, $\mathbf {\pi }(\mathbf {Z})_{ij} = \operatornamewithlimits{\mathsf {softmax}}(\mathbf {z}_i)_j$, where
In words, the keys are used to compute a relevance score between each item and query. Then, normalized attention weights are computed using softmax, and these are used to weight the values of each item at each query context.
However, for complex tasks, different parts of a sequence may be relevant in different ways, motivating multi-head attention in Transformers. This is simply the application of Equation DISPLAY_FORM7 in parallel $H$ times, each with a different, learned linear transformation that allows specialization:
In the Transformer, there are three separate multi-head attention mechanisms for distinct purposes:
Encoder self-attention: builds rich, layered representations of each input word, by attending on the entire input sentence.
Context attention: selects a representative weighted average of the encodings of the input words, at each time step of the decoder.
Decoder self-attention: attends over the partial output sentence fragment produced so far.
Together, these mechanisms enable the contextualized flow of information between the input sentence and the sequential decoder.
The softmax mapping (Equation DISPLAY_FORM8) is elementwise proportional to $\exp $, therefore it can never assign a weight of exactly zero. Thus, unnecessary items are still taken into consideration to some extent. Since its output sums to one, this invariably means less weight is assigned to the relevant items, potentially harming performance and interpretability BIBREF18. This has motivated a line of research on learning networks with sparse mappings BIBREF19, BIBREF20, BIBREF21, BIBREF22. We focus on a recently-introduced flexible family of transformations, $\alpha $-entmax BIBREF23, BIBREF14, defined as:
where $\triangle ^d \lbrace \mathbf {p}\in \mathbb {R}^d:\sum _{i} p_i = 1\rbrace $ is the probability simplex, and, for $\alpha \ge 1$, $\mathsf {H}^{\textsc {T}}_\alpha $ is the Tsallis continuous family of entropies BIBREF24:
This family contains the well-known Shannon and Gini entropies, corresponding to the cases $\alpha =1$ and $\alpha =2$, respectively.
Equation DISPLAY_FORM14 involves a convex optimization subproblem. Using the definition of $\mathsf {H}^{\textsc {T}}_\alpha $, the optimality conditions may be used to derive the following form for the solution (Appendix SECREF83):
where $[\cdot ]_+$ is the positive part (ReLU) function, $\mathbf {1}$ denotes the vector of all ones, and $\tau $ – which acts like a threshold – is the Lagrange multiplier corresponding to the $\sum _i p_i=1$ constraint.
The appeal of $\alpha $-entmax for attention rests on the following properties. For $\alpha =1$ (i.e., when $\mathsf {H}^{\textsc {T}}_\alpha $ becomes the Shannon entropy), it exactly recovers the softmax mapping (We provide a short derivation in Appendix SECREF89.). For all $\alpha >1$ it permits sparse solutions, in stark contrast to softmax. In particular, for $\alpha =2$, it recovers the sparsemax mapping BIBREF19, which is piecewise linear. In-between, as $\alpha $ increases, the mapping continuously gets sparser as its curvature changes.
To compute the value of $\alpha $-entmax, one must find the threshold $\tau $ such that the r.h.s. in Equation DISPLAY_FORM16 sums to one. BIBREF23 propose a general bisection algorithm. BIBREF14 introduce a faster, exact algorithm for $\alpha =1.5$, and enable using $\mathop {\mathsf {\alpha }\textnormal {-}\mathsf {entmax }}$ with fixed $\alpha $ within a neural network by showing that the $\alpha $-entmax Jacobian w.r.t. $\mathbf {z}$ for $\mathbf {p}^\star = \mathop {\mathsf {\alpha }\textnormal {-}\mathsf {entmax }}(\mathbf {z})$ is
Our work furthers the study of $\alpha $-entmax by providing a derivation of the Jacobian w.r.t. the hyper-parameter $\alpha $ (Section SECREF3), thereby allowing the shape and sparsity of the mapping to be learned automatically. This is particularly appealing in the context of multi-head attention mechanisms, where we shall show in Section SECREF35 that different heads tend to learn different sparsity behaviors.
We now propose a novel Transformer architecture wherein we simply replace softmax with $\alpha $-entmax in the attention heads. Concretely, we replace the row normalization $\mathbf {\pi }$ in Equation DISPLAY_FORM7 by
This change leads to sparse attention weights, as long as $\alpha >1$; in particular, $\alpha =1.5$ is a sensible starting point BIBREF14.
Unlike LSTM-based seq2seq models, where $\alpha $ can be more easily tuned by grid search, in a Transformer, there are many attention heads in multiple layers. Crucial to the power of such models, the different heads capture different linguistic phenomena, some of them isolating important words, others spreading out attention across phrases BIBREF0. This motivates using different, adaptive $\alpha $ values for each attention head, such that some heads may learn to be sparser, and others may become closer to softmax. We propose doing so by treating the $\alpha $ values as neural network parameters, optimized via stochastic gradients along with the other weights.
In order to optimize $\alpha $ automatically via gradient methods, we must compute the Jacobian of the entmax output w.r.t. $\alpha $. Since entmax is defined through an optimization problem, this is non-trivial and cannot be simply handled through automatic differentiation; it falls within the domain of argmin differentiation, an active research topic in optimization BIBREF25, BIBREF26.
One of our key contributions is the derivation of a closed-form expression for this Jacobian. The next proposition provides such an expression, enabling entmax layers with adaptive $\alpha $. To the best of our knowledge, ours is the first neural network module that can automatically, continuously vary in shape away from softmax and toward sparse mappings like sparsemax.
Proposition 1 Let $\mathbf {p}^\star \mathop {\mathsf {\alpha }\textnormal {-}\mathsf {entmax }}(\mathbf {z})$ be the solution of Equation DISPLAY_FORM14. Denote the distribution $\tilde{p}_i {(p_i^\star )^{2 - \alpha }}{ \sum _j(p_j^\star )^{2-\alpha }}$ and let $h_i -p^\star _i \log p^\star _i$. The $i$th component of the Jacobian $\mathbf {g} \frac{\partial \mathop {\mathsf {\alpha }\textnormal {-}\mathsf {entmax }}(\mathbf {z})}{\partial \alpha }$ is
proof uses implicit function differentiation and is given in Appendix SECREF10.
Proposition UNKREF22 provides the remaining missing piece needed for training adaptively sparse Transformers. In the following section, we evaluate this strategy on neural machine translation, and analyze the behavior of the learned attention heads.
We apply our adaptively sparse Transformers on four machine translation tasks. For comparison, a natural baseline is the standard Transformer architecture using the softmax transform in its multi-head attention mechanisms. We consider two other model variants in our experiments that make use of different normalizing transformations:
1.5-entmax: a Transformer with sparse entmax attention with fixed $\alpha =1.5$ for all heads. This is a novel model, since 1.5-entmax had only been proposed for RNN-based NMT models BIBREF14, but never in Transformers, where attention modules are not just one single component of the seq2seq model but rather an integral part of all of the model components.
$\alpha $-entmax: an adaptive Transformer with sparse entmax attention with a different, learned $\alpha _{i,j}^t$ for each head.
The adaptive model has an additional scalar parameter per attention head per layer for each of the three attention mechanisms (encoder self-attention, context attention, and decoder self-attention), i.e.,
and we set $\alpha _{i,j}^t = 1 + \operatornamewithlimits{\mathsf {sigmoid}}(a_{i,j}^t) \in ]1, 2[$. All or some of the $\alpha $ values can be tied if desired, but we keep them independent for analysis purposes.
Our models were trained on 4 machine translation datasets of different training sizes:
[itemsep=.5ex,leftmargin=2ex]
IWSLT 2017 German $\rightarrow $ English BIBREF27: 200K sentence pairs.
KFTT Japanese $\rightarrow $ English BIBREF28: 300K sentence pairs.
WMT 2016 Romanian $\rightarrow $ English BIBREF29: 600K sentence pairs.
WMT 2014 English $\rightarrow $ German BIBREF30: 4.5M sentence pairs.
All of these datasets were preprocessed with byte-pair encoding BIBREF31, using joint segmentations of 32k merge operations.
We follow the dimensions of the Transformer-Base model of BIBREF0: The number of layers is $L=6$ and number of heads is $H=8$ in the encoder self-attention, the context attention, and the decoder self-attention. We use a mini-batch size of 8192 tokens and warm up the learning rate linearly until 20k steps, after which it decays according to an inverse square root schedule. All models were trained until convergence of validation accuracy, and evaluation was done at each 10k steps for ro$\rightarrow $en and en$\rightarrow $de and at each 5k steps for de$\rightarrow $en and ja$\rightarrow $en. The end-to-end computational overhead of our methods, when compared to standard softmax, is relatively small; in training tokens per second, the models using $\alpha $-entmax and $1.5$-entmax are, respectively, $75\%$ and $90\%$ the speed of the softmax model.
We report test set tokenized BLEU BIBREF32 results in Table TABREF27. We can see that replacing softmax by entmax does not hurt performance in any of the datasets; indeed, sparse attention Transformers tend to have slightly higher BLEU, but their sparsity leads to a better potential for analysis. In the next section, we make use of this potential by exploring the learned internal mechanics of the self-attention heads.
We conduct an analysis for the higher-resource dataset WMT 2014 English $\rightarrow $ German of the attention in the sparse adaptive Transformer model ($\alpha $-entmax) at multiple levels: we analyze high-level statistics as well as individual head behavior. Moreover, we make a qualitative analysis of the interpretability capabilities of our models.
Figure FIGREF37 shows the learning trajectories of the $\alpha $ parameters of a selected subset of heads. We generally observe a tendency for the randomly-initialized $\alpha $ parameters to decrease initially, suggesting that softmax-like behavior may be preferable while the model is still very uncertain. After around one thousand steps, some heads change direction and become sparser, perhaps as they become more confident and specialized. This shows that the initialization of $\alpha $ does not predetermine its sparsity level or the role the head will have throughout. In particular, head 8 in the encoder self-attention layer 2 first drops to around $\alpha =1.3$ before becoming one of the sparsest heads, with $\alpha \approx 2$.
The overall distribution of $\alpha $ values at convergence can be seen in Figure FIGREF38. We can observe that the encoder self-attention blocks learn to concentrate the $\alpha $ values in two modes: a very sparse one around $\alpha \rightarrow 2$, and a dense one between softmax and 1.5-entmax . However, the decoder self and context attention only learn to distribute these parameters in a single mode. We show next that this is reflected in the average density of attention weight vectors as well.
For any $\alpha >1$, it would still be possible for the weight matrices in Equation DISPLAY_FORM9 to learn re-scalings so as to make attention sparser or denser. To visualize the impact of adaptive $\alpha $ values, we compare the empirical attention weight density (the average number of tokens receiving non-zero attention) within each module, against sparse Transformers with fixed $\alpha =1.5$.
Figure FIGREF40 shows that, with fixed $\alpha =1.5$, heads tend to be sparse and similarly-distributed in all three attention modules. With learned $\alpha $, there are two notable changes: (i) a prominent mode corresponding to fully dense probabilities, showing that our models learn to combine sparse and dense attention, and (ii) a distinction between the encoder self-attention – whose background distribution tends toward extreme sparsity – and the other two modules, who exhibit more uniform background distributions. This suggests that perhaps entirely sparse Transformers are suboptimal.
The fact that the decoder seems to prefer denser attention distributions might be attributed to it being auto-regressive, only having access to past tokens and not the full sentence. We speculate that it might lose too much information if it assigned weights of zero to too many tokens in the self-attention, since there are fewer tokens to attend to in the first place.
Teasing this down into separate layers, Figure FIGREF41 shows the average (sorted) density of each head for each layer. We observe that $\alpha $-entmax is able to learn different sparsity patterns at each layer, leading to more variance in individual head behavior, to clearly-identified dense and sparse heads, and overall to different tendencies compared to the fixed case of $\alpha =1.5$.
To measure the overall disagreement between attention heads, as a measure of head diversity, we use the following generalization of the Jensen-Shannon divergence:
where $\mathbf {p}_j$ is the vector of attention weights assigned by head $j$ to each word in the sequence, and $\mathsf {H}^\textsc {S}$ is the Shannon entropy, base-adjusted based on the dimension of $\mathbf {p}$ such that $JS \le 1$. We average this measure over the entire validation set. The higher this metric is, the more the heads are taking different roles in the model.
Figure FIGREF44 shows that both sparse Transformer variants show more diversity than the traditional softmax one. Interestingly, diversity seems to peak in the middle layers of the encoder self-attention and context attention, while this is not the case for the decoder self-attention.
The statistics shown in this section can be found for the other language pairs in Appendix SECREF8.
Previous work pointed out some specific roles played by different heads in the softmax Transformer model BIBREF33, BIBREF5, BIBREF9. Identifying the specialization of a head can be done by observing the type of tokens or sequences that the head often assigns most of its attention weight; this is facilitated by sparsity.
One particular type of head, as noted by BIBREF9, is the positional head. These heads tend to focus their attention on either the previous or next token in the sequence, thus obtaining representations of the neighborhood of the current time step. In Figure FIGREF47, we show attention plots for such heads, found for each of the studied models. The sparsity of our models allows these heads to be more confident in their representations, by assigning the whole probability distribution to a single token in the sequence. Concretely, we may measure a positional head's confidence as the average attention weight assigned to the previous token. The softmax model has three heads for position $-1$, with median confidence $93.5\%$. The $1.5$-entmax model also has three heads for this position, with median confidence $94.4\%$. The adaptive model has four heads, with median confidences $95.9\%$, the lowest-confidence head being dense with $\alpha =1.18$, while the highest-confidence head being sparse ($\alpha =1.91$).
For position $+1$, the models each dedicate one head, with confidence around $95\%$, slightly higher for entmax. The adaptive model sets $\alpha =1.96$ for this head.
Due to the sparsity of our models, we are able to identify other head specializations, easily identifying which heads should be further analysed. In Figure FIGREF51 we show one such head where the $\alpha $ value is particularly high (in the encoder, layer 1, head 4 depicted in Figure FIGREF37). We found that this head most often looks at the current time step with high confidence, making it a positional head with offset 0. However, this head often spreads weight sparsely over 2-3 neighboring tokens, when the tokens are part of the same BPE cluster or hyphenated words. As this head is in the first layer, it provides a useful service to the higher layers by combining information evenly within some BPE clusters.
For each BPE cluster or cluster of hyphenated words, we computed a score between 0 and 1 that corresponds to the maximum attention mass assigned by any token to the rest of the tokens inside the cluster in order to quantify the BPE-merging capabilities of these heads. There are not any attention heads in the softmax model that are able to obtain a score over $80\%$, while for $1.5$-entmax and $\alpha $-entmax there are two heads in each ($83.3\%$ and $85.6\%$ for $1.5$-entmax and $88.5\%$ and $89.8\%$ for $\alpha $-entmax).
On the other hand, in Figure FIGREF52 we show a head for which our adaptively sparse model chose an $\alpha $ close to 1, making it closer to softmax (also shown in encoder, layer 1, head 3 depicted in Figure FIGREF37). We observe that this head assigns a high probability to question marks at the end of the sentence in time steps where the current token is interrogative, thus making it an interrogation-detecting head. We also observe this type of heads in the other models, which we also depict in Figure FIGREF52. The average attention weight placed on the question mark when the current token is an interrogative word is $98.5\%$ for softmax, $97.0\%$ for $1.5$-entmax, and $99.5\%$ for $\alpha $-entmax.
Furthermore, we can examine sentences where some tendentially sparse heads become less so, thus identifying sources of ambiguity where the head is less confident in its prediction. An example is shown in Figure FIGREF55 where sparsity in the same head differs for sentences of similar length.
Prior work has developed sparse attention mechanisms, including applications to NMT BIBREF19, BIBREF12, BIBREF20, BIBREF22, BIBREF34. BIBREF14 introduced the entmax function this work builds upon. In their work, there is a single attention mechanism which is controlled by a fixed $\alpha $. In contrast, this is the first work to allow such attention mappings to dynamically adapt their curvature and sparsity, by automatically adjusting the continuous $\alpha $ parameter. We also provide the first results using sparse attention in a Transformer model.
Recent research improves the scalability of Transformer-like networks through static, fixed sparsity patterns BIBREF10, BIBREF35. Our adaptively-sparse Transformer can dynamically select a sparsity pattern that finds relevant words regardless of their position (e.g., Figure FIGREF52). Moreover, the two strategies could be combined. In a concurrent line of research, BIBREF11 propose an adaptive attention span for Transformer language models. While their work has each head learn a different contiguous span of context tokens to attend to, our work finds different sparsity patterns in the same span. Interestingly, some of their findings mirror ours – we found that attention heads in the last layers tend to be denser on average when compared to the ones in the first layers, while their work has found that lower layers tend to have a shorter attention span compared to higher layers.
The original Transformer paper BIBREF0 shows attention visualizations, from which some speculation can be made of the roles the several attention heads have. BIBREF7 study the syntactic abilities of the Transformer self-attention, while BIBREF6 extract dependency relations from the attention weights. BIBREF8 find that the self-attentions in BERT BIBREF3 follow a sequence of processes that resembles a classical NLP pipeline. Regarding redundancy of heads, BIBREF9 develop a method that is able to prune heads of the multi-head attention module and make an empirical study of the role that each head has in self-attention (positional, syntactic and rare words). BIBREF36 also aim to reduce head redundancy by adding a regularization term to the loss that maximizes head disagreement and obtain improved results. While not considering Transformer attentions, BIBREF18 show that traditional attention mechanisms do not necessarily improve interpretability since softmax attention is vulnerable to an adversarial attack leading to wildly different model predictions for the same attention weights. Sparse attention may mitigate these issues; however, our work focuses mostly on a more mechanical aspect of interpretation by analyzing head behavior, rather than on explanations for predictions.
We contribute a novel strategy for adaptively sparse attention, and, in particular, for adaptively sparse Transformers. We present the first empirical analysis of Transformers with sparse attention mappings (i.e., entmax), showing potential in both translation accuracy as well as in model interpretability.
In particular, we analyzed how the attention heads in the proposed adaptively sparse Transformer can specialize more and with higher confidence. Our adaptivity strategy relies only on gradient-based optimization, side-stepping costly per-head hyper-parameter searches. Further speed-ups are possible by leveraging more parallelism in the bisection algorithm for computing $\alpha $-entmax.
Finally, some of the automatically-learned behaviors of our adaptively sparse Transformers – for instance, the near-deterministic positional heads or the subword joining head – may provide new ideas for designing static variations of the Transformer.
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 are grateful to Ben Peters for the $\alpha $-entmax code and Erick Fonseca, Marcos Treviso, Pedro Martins, and Tsvetomila Mihaylova for insightful group discussion. We thank Mathieu Blondel for the idea to learn $\alpha $. We would also like to thank the anonymous reviewers for their helpful feedback.
Supplementary Material
Definition 1 (BIBREF23)
Let $\Omega \colon \triangle ^d \rightarrow {\mathbb {R}}\cup \lbrace \infty \rbrace $ be a strictly convex regularization function. We define the prediction function $\mathbf {\pi }_{\Omega }$ as
Lemma 1 (BIBREF14) For any $\mathbf {z}$, there exists a unique $\tau ^\star $ such that
Proof: From the definition of $\mathop {\mathsf {\alpha }\textnormal {-}\mathsf {entmax }}$,
we may easily identify it with a regularized prediction function (Def. UNKREF81):
We first note that for all $\mathbf {p}\in \triangle ^d$,
From the constant invariance and scaling properties of $\mathbf {\pi }_{\Omega }$ BIBREF23,
Using BIBREF23, noting that $g^{\prime }(t) = t^{\alpha - 1}$ and $(g^{\prime })^{-1}(u) = u^{{1}{\alpha -1}}$, yields
Since $\mathsf {H}^{\textsc {T}}_\alpha $ is strictly convex on the simplex, $\mathop {\mathsf {\alpha }\textnormal {-}\mathsf {entmax }}$ has a unique solution $\mathbf {p}^\star $. Equation DISPLAY_FORM88 implicitly defines a one-to-one mapping between $\mathbf {p}^\star $ and $\tau ^\star $ as long as $\mathbf {p}^\star \in \triangle $, therefore $\tau ^\star $ is also unique.
The Euclidean projection onto the simplex, sometimes referred to, in the context of neural attention, as sparsemax BIBREF19, is defined as
The solution can be characterized through the unique threshold $\tau $ such that $\sum _i \operatornamewithlimits{\mathsf {sparsemax}}(\mathbf {z})_i = 1$ and BIBREF38
Thus, each coordinate of the sparsemax solution is a piecewise-linear function. Visibly, this expression is recovered when setting $\alpha =2$ in the $\alpha $-entmax expression (Equation DISPLAY_FORM85); for other values of $\alpha $, the exponent induces curvature.
On the other hand, the well-known softmax is usually defined through the expression
which can be shown to be the unique solution of the optimization problem
where $\mathsf {H}^\textsc {S}(\mathbf {p}) -\sum _i p_i \log p_i$ is the Shannon entropy. Indeed, setting the gradient to 0 yields the condition $\log p_i = z_j - \nu _i - \tau - 1$, where $\tau $ and $\nu > 0$ are Lagrange multipliers for the simplex constraints $\sum _i p_i = 1$ and $p_i \ge 0$, respectively. Since the l.h.s. is only finite for $p_i>0$, we must have $\nu _i=0$ for all $i$, by complementary slackness. Thus, the solution must have the form $p_i = {\exp (z_i)}{Z}$, yielding Equation DISPLAY_FORM92.
Recall that the entmax transformation is defined as:
where $\alpha \ge 1$ and $\mathsf {H}^{\textsc {T}}_{\alpha }$ is the Tsallis entropy,
and $\mathsf {H}^\textsc {S}(\mathbf {p}):= -\sum _j p_j \log p_j$ is the Shannon entropy.
In this section, we derive the Jacobian of $\operatornamewithlimits{\mathsf {entmax }}$ with respect to the scalar parameter $\alpha $.
From the KKT conditions associated with the optimization problem in Eq. DISPLAY_FORM85, we have that the solution $\mathbf {p}^{\star }$ has the following form, coordinate-wise:
where $\tau ^{\star }$ is a scalar Lagrange multiplier that ensures that $\mathbf {p}^{\star }$ normalizes to 1, i.e., it is defined implicitly by the condition:
For general values of $\alpha $, Eq. DISPLAY_FORM98 lacks a closed form solution. This makes the computation of the Jacobian
non-trivial. Fortunately, we can use the technique of implicit differentiation to obtain this Jacobian.
The Jacobian exists almost everywhere, and the expressions we derive expressions yield a generalized Jacobian BIBREF37 at any non-differentiable points that may occur for certain ($\alpha $, $\mathbf {z}$) pairs. We begin by noting that $\frac{\partial p_i^{\star }}{\partial \alpha } = 0$ if $p_i^{\star } = 0$, because increasing $\alpha $ keeps sparse coordinates sparse. Therefore we need to worry only about coordinates that are in the support of $\mathbf {p}^\star $. We will assume hereafter that the $i$th coordinate of $\mathbf {p}^\star $ is non-zero. We have:
We can see that this Jacobian depends on $\frac{\partial \tau ^{\star }}{\partial \alpha }$, which we now compute using implicit differentiation.
Let $\mathcal {S} = \lbrace i: p^\star _i > 0 \rbrace $). By differentiating both sides of Eq. DISPLAY_FORM98, re-using some of the steps in Eq. DISPLAY_FORM101, and recalling Eq. DISPLAY_FORM97, we get
from which we obtain:
Finally, plugging Eq. DISPLAY_FORM103 into Eq. DISPLAY_FORM101, we get:
where we denote by
The distribution $\tilde{\mathbf {p}}(\alpha )$ can be interpreted as a “skewed” distribution obtained from $\mathbf {p}^{\star }$, which appears in the Jacobian of $\mathop {\mathsf {\alpha }\textnormal {-}\mathsf {entmax }}(\mathbf {z})$ w.r.t. $\mathbf {z}$ as well BIBREF14.
We can write Eq. DISPLAY_FORM104 as
When $\alpha \rightarrow 1^+$, we have $\tilde{\mathbf {p}}(\alpha ) \rightarrow \mathbf {p}^{\star }$, which leads to a $\frac{0}{0}$ indetermination.
To solve this indetermination, we will need to apply L'Hôpital's rule twice. Let us first compute the derivative of $\tilde{p}_i(\alpha )$ with respect to $\alpha $. We have
therefore
Differentiating the numerator and denominator in Eq. DISPLAY_FORM107, we get:
with
and
When $\alpha \rightarrow 1^+$, $B$ becomes again a $\frac{0}{0}$ indetermination, which we can solve by applying again L'Hôpital's rule. Differentiating the numerator and denominator in Eq. DISPLAY_FORM112:
Finally, summing Eq. DISPLAY_FORM111 and Eq. DISPLAY_FORM113, we get
To sum up, we have the following expression for the Jacobian of $\mathop {\mathsf {\alpha }\textnormal {-}\mathsf {entmax }}$ with respect to $\alpha $:

Figure and table captions:
Figure 1: Attention distributions of different selfattention heads for the time step of the token “over”, shown to compare our model to other related work. While the sparse Transformer (Child et al., 2019) and the adaptive span Transformer (Sukhbaatar et al., 2019) only attend to words within a contiguous span of the past tokens, our model is not only able to obtain different and not necessarily contiguous sparsity patterns for each attention head, but is also able to tune its support over which tokens to attend adaptively.
Table 1: Machine translation tokenized BLEU test results on IWSLT 2017 DE EN, KFTT JA EN, WMT 2016 RO EN and WMT 2014 EN DE, respectively.
Figure 2: Trajectories of α values for a subset of the heads during training. Initialized at random, most heads become denser in the beginning, before converging. This suggests that dense attention may be more beneficial while the network is still uncertain, being replaced by sparse attention afterwards.
Figure 4: Distribution of attention densities (average number of tokens receiving non-zero attention weight) for all attention heads and all validation sentences. When compared to 1.5-entmax, α-entmax distributes the sparsity in a more uniform manner, with a clear mode at fully dense attentions, corresponding to the heads with low α. In the softmax case, this distribution would lead to a single bar with density 1.
Figure 3: Distribution of learned α values per attention block. While the encoder self-attention has a bimodal distribution of values of α, the decoder self-attention and context attention have a single mode.
Figure 5: Head density per layer for fixed and learned α. Each line corresponds to an attention head; lower values mean that that attention head is sparser. Learned α has higher variance.
Figure 6: Jensen-Shannon Divergence between heads at each layer. Measures the disagreement between heads: the higher the value, the more the heads are disagreeing with each other in terms of where to attend. Models using sparse entmax have more diverse attention than the softmax baseline.
Figure 9: Interrogation-detecting heads in the three models. The top sentence is interrogative while the bottom one is declarative but includes the interrogative word “what”. In the top example, these interrogation heads assign a high probability to the question mark in the time step of the interrogative word (with ≥ 97.0% probability), while in the bottom example since there is no question mark, the same head does not assign a high probability to the last token in the sentence during the interrogative word time step. Surprisingly, this head prefers a low α = 1.05, as can be seen from the dense weights. This allows the head to identify the noun phrase “Armani Polo” better.
Figure 7: Self-attention from the most confidently previous-position head in each model. The learned parameter in the α-entmax model is α = 1.91. Quantitatively more confident, visual inspection confirms that the adaptive head behaves more consistently.
Figure 8: BPE-merging head (α = 1.91) discovered in the α-entmax model. Found in the first encoder layer, this head learns to discover some subword units and combine their information, leaving most words intact. It places 99.09% of its probability mass within the same BPE cluster as the current token: more than any head in any other model.
Figure 10: Example of two sentences of similar length where the same head (α = 1.33) exhibits different sparsity. The longer phrase in the example on the right “a sexually transmitted disease” is handled with higher confidence, leading to more sparsity.
Figure 11: Histograms of α values.
Figure 12: Histograms of head densities.
Figure 13: Jensen-Shannon divergence over layers.
Figure 14: Head densities over layers.

Questions about this paper:
What tasks are used for evaluation?
HOw does the method perform compared with baselines?

Answers:
four machine translation tasks: German -> English, Japanese -> English, Romanian -> English, English -> German
On the datasets DE-EN, JA-EN, RO-EN, and EN-DE, the baseline achieves 29.79, 21.57, 32.70, and 26.02  BLEU score, respectively. The 1.5-entmax achieves  29.83, 22.13, 33.10, and 25.89 BLEU score, which is a difference of +0.04, +0.56, +0.40, and -0.13 BLEU score versus the baseline. The α-entmax achieves 29.90, 21.74, 32.89, and 26.93 BLEU score, which is a difference of +0.11, +0.17, +0.19, +0.91 BLEU score versus the baseline.