Sparse and Constrained Attention for Neural Machine Translation

Introduction

Neural machine translation (NMT) emerged in the last few years as a very successful paradigm (Sutskever et al., 2014; Bahdanau et al., 2014; Gehring et al., 2017; Vaswani et al., 2017). While NMT is generally more fluent than previous statistical systems, adequacy is still a major concern Koehn and Knowles (2017): common mistakes include dropping source words and repeating words in the generated translation.

Previous work has attempted to mitigate this problem in various ways. Wu et al. (2016) incorporate coverage and length penalties during beam search—a simple yet limited solution, since it only affects the scores of translation hypotheses that are already in the beam. Other approaches involve architectural changes: providing coverage vectors to track the attention history Mi et al. (2016); Tu et al. (2016), using gating architectures and adaptive attention to control the amount of source context provided Tu et al. (2017a); Li and Zhu (2017), or adding a reconstruction loss Tu et al. (2017b). Feng et al. (2016) also use the notion of fertility implicitly in their proposed model. Their “fertility conditioned decoder” uses a coverage vector and an “extract gate” which are incorporated in the decoding recurrent unit, increasing the number of parameters.

In this paper, we propose a different solution that does not change the overall architecture, but only the attention transformation. Namely, we replace the traditional softmax by other recently proposed transformations that either promote attention sparsity (Martins and Astudillo, 2016) or upper bound the amount of attention a word can receive (Martins and Kreutzer, 2017). The bounds are determined by the fertility values of the source words. While these transformations have given encouraging results in various NLP problems, they have never been applied to NMT, to the best of our knowledge. Furthermore, we combine these two ideas and propose a novel attention transformation, constrained sparsemax, which produces both sparse and bounded attention weights, yielding a compact and interpretable set of alignments. While being in-between soft and hard alignments (Figure 2), the constrained sparsemax transformation is end-to-end differentiable, hence amenable for training with gradient backpropagation.

To sum up, our contributions are as follows:Our software code is available at the OpenNMT fork www.github.com/Unbabel/OpenNMT-py/tree/dev and the running scripts at www.github.com/Unbabel/ sparse_constrained_attention.

We formulate constrained sparsemax and derive efficient linear and sublinear-time algorithms for running forward and backward propagation. This transformation has two levels of sparsity: over time steps, and over the attended words at each step.

We provide a detailed empirical comparison of various attention transformations, including softmax (Bahdanau et al., 2014), sparsemax (Martins and Astudillo, 2016), constrained softmax (Martins and Kreutzer, 2017), and our newly proposed constrained sparsemax. We provide error analysis including two new metrics targeted at detecting coverage problems.

Preliminaries

where $p(y_{t}\,\,|\,\,y_{1:(t-1)},x)$ is computed by a softmax output layer that receives a decoder state $\bm{s}_{t}$ as input. This state is updated by an auto-regressive LSTM, $\bm{s}_{t}=\operatorname*{\mathsf{RNN}}(\operatorname*{\mathsf{embed}}(y_{t-1}),\bm{s}_{t-1},\bm{c}_{t})$, where $\bm{c}_{t}$ is an input context vector. This vector is computed as $\bm{c}_{t}:=\bm{H}\bm{\alpha}_{t}$, where $\bm{\alpha}_{t}$ is a probability distribution that represents the attention over the source words, commonly obtained as

Sparse and Constrained Attention

In this work, we consider alternatives to Eq. 2. Since the softmax is strictly positive, it forces all words in the source to receive some probability mass in the resulting attention distribution, which can be wasteful. Moreover, it may happen that the decoder attends repeatedly to the same source words across time steps, causing repetitions in the generated translation, as Tu et al. (2016) observed.

The sparsemax transformation (Martins and Astudillo, 2016) is defined as:

Constrained softmax.

The constrained softmax transformation was recently proposed by Martins and Kreutzer (2017) in the context of easy-first sequence tagging, being defined as follows:

where $\bm{u}$ is a vector of upper bounds, and $\mathsf{KL}(.\|.)$ is the Kullback-Leibler divergence. In other words, it returns the distribution closest to $\operatorname*{\mathsf{softmax}}(\bm{z})$ whose attention probabilities are bounded by $\bm{u}$. Martins and Kreutzer (2017) have shown that this transformation can be evaluated in $O(J\log J)$ time and its gradients backpropagated in $O(J)$ time.

To use this transformation in the attention mechanism, we make use of the idea of fertility (Brown et al., 1993). Namely, let $\bm{\beta}_{t-1}:=\sum_{\tau=1}^{t-1}\bm{\alpha}_{\tau}$ denote the cumulative attention that each source word has received up to time step $t$, and let $\bm{f}:=(f_{j})_{j=1}^{J}$ be a vector containing fertility upper bounds for each source word. The attention at step $t$ is computed as

Intuitively, each source word $j$ gets a credit of $f_{j}$ units of attention, which are consumed along the decoding process. If all the credit is exhausted, it receives zero attention from then on. Unlike the sparsemax transformation, which places sparse attention over the source words, the constrained softmax leads to sparsity over time steps.

Constrained sparsemax.

In this work, we propose a novel transformation which shares the two properties above: it provides both sparse and bounded probabilities. It is defined as:

The following result, whose detailed proof we include as supplementary material (Appendix A), is key for enabling the use of the constrained sparsemax transformation in neural networks.

Let $\bm{\alpha}^{\star}=\operatorname*{\mathsf{csparsemax}}(\bm{z};\bm{u})$ be the solution of Eq. 6, and define the sets $\mathcal{A}=\{j\in[J]\,\,|\,\,0<\alpha_{j}^{\star}<u_{j}\}$, $\mathcal{A}_{L}=\{j\in[J]\,\,|\,\,\alpha_{j}^{\star}=0\}$, and $\mathcal{A}_{R}=\{j\in[J]\,\,|\,\,\alpha_{j}^{\star}=u_{j}\}$. Then:

Forward propagation. $\bm{\alpha}^{\star}$ can be computed in $O(J)$ time with the algorithm of Pardalos and Kovoor (1990) (Alg. 1 in Appendix A). The solution takes the form $\alpha_{j}^{\star}=\max\{0,\min\{u_{j},z_{j}-\tau\}\}$, where $\tau$ is a normalization constant.

Fertility Bounds

We experiment with three ways of setting the fertility of the source words: constant, guided, and predicted. With constant, we set the fertilities of all source words to a fixed integer value $f$. With guided, we train a word aligner based on IBM Model 2 (we used fast_align in our experiments, Dyer et al. (2013)) and, for each word in the vocabulary, we set the fertilities to the maximal observed value in the training data (or 1 if no alignment was observed). With the predicted strategy, we train a separate fertility predictor model using a bi-LSTM tagger.A similar strategy was recently used by Gu et al. (2018) as a component of their non-autoregressive NMT model. At training time, we provide as supervision the fertility estimated by fast_align. Since our model works with fertility upper bounds and the word aligner may miss some word pairs, we found it beneficial to add a constant to this number (1 in our experiments). At test time, we use the expected fertilities according to our model.

We append an additional $<$sink$>$ token to the end of the source sentence, to which we assign unbounded fertility ($f_{J+1}=\infty$). The token is akin to the null alignment in IBM models. The reason we add this token is the following: without the sink token, the length of the generated target sentence can never exceed $\sum_{j}f_{j}$ words if we use constrained softmax/sparsemax. At training time this may be problematic, since the target length is fixed and the problems in Eqs. 4–6 can become infeasible. By adding the sink token we guarantee $\sum_{j}f_{j}=\infty$, eliminating the problem.

Exhaustion strategies.

To avoid missing source words, we implemented a simple strategy to encourage more attention to words with larger credit: we redefine the pre-attention word scores as $\bm{z}_{t}^{\prime}=\bm{z}_{t}+c\bm{u}_{t}$, where $c$ is a constant ($c=0.2$ in our experiments). This increases the score of words which have not yet exhausted their fertility (we may regard it as a “soft” lower bound in Eqs. 4–6).

Experiments

We evaluated our attention transformations on three language pairs. We focused on small datasets, as they are the most affected by coverage mistakes. We use the IWSLT 2014 corpus for De-En, the KFTT corpus for Ja-En (Neubig, 2011), and the WMT 2016 dataset for Ro-En. The training sets have 153,326, 329,882, and 560,767 parallel sentences, respectively. Our reason to prefer smaller datasets is that this regime is what brings more adequacy issues and demands more structural biases, hence it is a good test bed for our methods. We tokenized the data using the Moses scripts and preprocessed it with subword units Sennrich et al. (2016) with a joint vocabulary and 32k merge operations. Our implementation was done on a fork of the OpenNMT-py toolkit Klein et al. (2017) with the default parameters We used a 2-layer LSTM, embedding and hidden size of 500, dropout 0.3, and the SGD optimizer for 13 epochs.. We used a validation set to tune hyperparameters introduced by our model. Even though our attention implementations are CPU-based using NumPy (unlike the rest of the computation which is done on the GPU), we did not observe any noticeable slowdown using multiple devices.

As baselines, we use softmax attention, as well as two recently proposed coverage models:

CovPenalty (Wu et al., 2016, §7). At test time, the hypotheses in the beam are rescored with a global score that includes a length and a coverage penalty.Since our sparse attention can become for some words, we extended the original coverage penalty by adding another parameter $\epsilon$, set to $0.1$: $\mathsf{cp}(x;y):=\beta\sum_{j=1}^{J}\log\max\{\epsilon,\min\{1,\sum_{t=1}^{|y|}\alpha_{jt}\}\}$. We tuned $\alpha$ and $\beta$ with grid search on $\{0.2k\}_{k=0}^{5}$, as in Wu et al. (2016).

CovVector (Tu et al., 2016). At training and test time, coverage vectors $\bm{\beta}$ and additional parameters $\bm{v}$ are used to condition the next attention step. We adapted this to our bilinear attention by defining $z_{t,j}=\bm{s}_{t-1}^{\top}(\bm{W}\bm{h}_{j}+\bm{v}{\beta}_{t-1,j})$.

We also experimented combining the strategies above with the sparsemax transformation.

As evaluation metrics, we report tokenized BLEU, METEOR (Denkowski and Lavie (2014), as well as two new metrics that we describe next to account for over and under-translation.Both evaluation metrics are included in our software package at www.github.com/Unbabel/ sparse_constrained_attention.

a new metric to count repetitions. Formally, given an $n$-gram $s\in V^{n}$, let $t(s)$ and $r(s)$ be the its frequency in the model translation and reference. We first compute a sentence-level score

The REP-score is then given by summing $\sigma(t,r)$ over sentences, normalizing by the number of words on the reference corpus, and multiplying by 100. We used $n=2$, $\lambda_{1}=1$ and $\lambda_{2}=2$.

DROP-score:

a new metric that accounts for possibly dropped words. To compute it, we first compute two sets of word alignments: from source to reference translation, and from source to the predicted translation. In our experiments, the alignments were obtained with fast_align Dyer et al. (2013), trained on the training partition of the data. Then, the DROP-score computes the percentage of source words that aligned with some word from the reference translation, but not with any word from the predicted translation.

Table 1 shows the results. We can see that on average, the sparse models ($\operatorname*{\mathsf{csparsemax}}$ as well as $\operatorname*{\mathsf{sparsemax}}$ combined with coverage models) have higher scores on both BLEU and METEOR. Generally, they also obtain better REP and DROP scores than $\operatorname*{\mathsf{csoftmax}}$ and $\operatorname*{\mathsf{softmax}}$, which suggests that sparse attention alleviates the problem of coverage to some extent.

To compare different fertility strategies, we ran experiments on the De-En for the $\operatorname*{\mathsf{csparsemax}}$ transformation (Table 2). We see that the Predicted strategy outperforms the others both in terms of BLEU and METEOR, albeit slightly.

Figure 2 shows examples of sentences for which the $\operatorname*{\mathsf{csparsemax}}$ fixed repetitions, along with the corresponding attention maps. We see that in the case of $\operatorname*{\mathsf{softmax}}$ repetitions, the decoder attends repeatedly to the same portion of the source sentence (the expression “letzten hundert” in the first sentence and “regierung” in the second sentence). Not only did $\operatorname*{\mathsf{csparsemax}}$ avoid repetitions, but it also yielded a sparse set of alignments, as expected. Appendix B provides more examples of translations from all models in discussion.

Conclusions

We proposed a new approach to address the coverage problem in NMT, by replacing the softmax attentional transformation by sparse and constrained alternatives: sparsemax, constrained softmax, and the newly proposed constrained sparsemax. For the latter, we derived efficient forward and backward propagation algorithms. By incorporating a model for fertility prediction, our attention transformations led to sparse alignments, avoiding repeated words in the translation.

Acknowledgments

We thank the Unbabel AI Research team for numerous discussions, and the three anonymous reviewers for their insightful comments. 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/2013, PTDC/EEI-SII/7092/2014 (LearnBig), and CMUPERI/TIC/0046/2014 (GoLocal).

References

Appendix A Proof of Proposition 1

We provide here a detailed proof of Proposition 1.

The optimization problem can be written as

To obtain the solution, we invoke the Karush-Kuhn-Tucker conditions. From the stationarity condition, we have $\mathbf{0}=\bm{\alpha}-\bm{z}+\tau\mathbf{1}-\bm{\mu}+\bm{\nu}$, which due to the primal feasibility condition implies that the solution is of the form:

From the complementarity slackness condition, we have that $0<\alpha_{j}<u_{j}$ implies that $\mu_{j}=\nu_{j}=0$ and therefore $\alpha_{j}=z_{j}-\tau$. On the other hand, $\mu_{j}>0$ implies $\alpha_{j}=0$, and $\nu_{j}>0$ implies $\alpha_{j}=u_{j}$. Hence the solution can be written as $\alpha_{j}=\max\{0,\min\{u_{j},z_{j}-\tau\}$, where $\tau$ is determined such that the distribution normalizes:

with $\mathcal{A}=\{j\in[J]\,\,|\,\,0<\alpha_{j}<u_{j}\}$ and $\mathcal{A}_{R}=\{j\in[J]\,\,|\,\,\alpha_{j}=u_{j}\}$. Note that $\tau$ depends itself on the set $\mathcal{A}$, a function of the solution. In §A.3, we describe an algorithm that searches the value of $\tau$ efficiently.

A.2 Gradient Backpropagation

We now turn to the problem of backpropagating the gradients through the constrained sparsemax transformation. For that, we need to compute its Jacobian matrix, i.e., the derivatives $\frac{\partial\alpha_{i}}{\partial z_{j}}$ and $\frac{\partial\alpha_{i}}{\partial u_{j}}$ for $i,j\in[J]$. Let us first express $\bm{\alpha}$ as

with $\tau$ as in Eq. 13. Note that we have ${\partial\tau}/{\partial z_{j}}=\mathds{1}(j\in\mathcal{A})/|\mathcal{A}|$ and ${\partial\tau}/{\partial u_{j}}=\mathds{1}(j\in\mathcal{A}_{R})/|\mathcal{A}|$. Thus, we have the following:

A.3 Linear-Time Evaluation

Finally, we present an algorithm to solve the problem in Eq. 6 in linear time.

Pardalos and Kovoor (1990) describe an algorithm, reproduced here as Algorithm 1, for solving a class of singly-constrained convex quadratic problems, which can be written in the form above (where each $c_{j}\geq 0$):

The solution of the problem in Eq. A.3 is of the form $x_{j}^{\star}=\max\{a_{j},\min\{b_{j},y\}\}$, where $y\in[a_{j},b_{j}]$ is a constant. The algorithm searches the value of this constant (which is similar to $\tau$ in our problem), which lies in a particular interval of split-points (line 3), iteratively shrinking this interval. The algorithm requires computing medians as a subroutine, which can be done in linear time (Blum et al., 1973). The overall complexity in $O(J)$ (Pardalos and Kovoor, 1990). The same algorithm has been used in NLP by Almeida and Martins (2013) for a budgeted summarization problem.

To show that this algorithm applies to the problem of evaluating $\operatorname*{\mathsf{csparsemax}}$, it suffices to show that our problem in Eq. 6 can be rewritten in the form of Eq. A.3. This is indeed the case, if we set:

Appendix B Examples of Translations

We show some examples of translations obtained for the German-English language pair with different systems. Blue highlights the parts of the reference that are correct and red highlights the corresponding problematic parts of translations, including repetitions, dropped words or mistranslations.