Gaussian Multi-head Attention for Simultaneous Machine Translation

In a translation task, we denote the source sentence as $\mathbf{x}\!=\!\left\{x_{1},\cdots,x_{J}\right\}$ and the corresponding source hidden states as $\mathbf{z}\!=\!\left\{z_{1},\cdots,z_{J}\right\}$ with source length $J$. The model generates a target sentence $\mathbf{y}\!=\!\left\{y_{1},\cdots,y_{I}\right\}$ and the corresponding target hidden states $\mathbf{s}\!=\!\left\{s_{1},\cdots,s_{I}\right\}$ with target length $I$. Different from the full-sentence machine translation, the source words received by SiMT model are incremental and hence the model needs to decide where to output translation.

Output position Define $g(i)$ Ma et al. (2019) as a monotonic non-decreasing function of step $i$, to denote the number of source words received by SiMT model when translating $y_{i}$, i.e., $g(i)$ is the output position of $y_{i}$.

In SiMT, $g(i)$ is determined by the specific policy, and the probability of generating the target word $y_{i}$ is $p\left(y_{i}\mid\mathbf{x}_{\leq g(i)},\mathbf{y}_{<i}\right)$, where $\mathbf{x}_{\leq g(i)}$ is first $g(i)$ source words and $\mathbf{y}_{<i}$ is previous target words. Therefore, the decoding probability of $\mathbf{y}$ is calculated as:

Multi-head attention Vaswani et al. (2017) contains multiple attention heads, where each attention head performs scaled dot-product attention. Our method is based on the cross-attention, where the queries are the target hidden states $\mathbf{s}$, the keys and values both come from the source hidden states $\mathbf{z}$. The soft attention weight $\alpha_{ij}$ is calculated as:

where $Q\!\left(\cdot\right)$ and $K\!\left(\cdot\right)$ are projection functions from the input space to the query and key space respectively, and $d_{k}$ is the dimension of inputs. Then the context vector $c_{i}$ is calculated as:

where $V\!\left(\cdot\right)$ is a projection function to value space.

Alignments prediction To develop a SiMT policy with alignments, GMA first predicts the aligned source position of the current target word. Due to the incrementality of streaming inputs in SiMT, it is unstable to directly predict the absolute position of the aligned source word. Instead, we predict the relative distance from the previous aligned source position, called incremental step.

Formally, we denote the aligned source position of the $i^{th}$ target word as $p_{i}\!\in\!\left[1,J\right]$ and the incremental step as $\Delta p_{i}\!\in\!\left(0,+\infty\right)$. Therefore, the aligned source position $p_{i}$ is calculated as:

where we set the initial aligned position $p_{0}$ to the first source word, and the incremental step $\Delta p_{i}$ is predicted through a multi-layer perceptron (MLP) based on the previous target hidden state $s_{i-1}$:

where $\mathbf{W_{\!p}}$, $\mathbf{V_{\!p}}$ are learnable parameters of MLP.

SiMT policy Besides the predicted aligned position $p_{i}$, we also introduce a relaxation offset $\delta$ to allow the model to wait for some additional source inputs, thereby providing a controllable trade-off between translation quality and latency in practice. Specifically, the output position $g(i)$ (i.e., wait for the first $g(i)$ source words and then translate the $i^{th}$ target word) is calculated as:

where $\left\lfloor\cdot\right\rfloor$ is a floor operation. In our experiments, relaxation offset $\delta$ is a hyperparameter we set to obtain the translation quality under different latency. Overall, the SiMT policy is shown in Algorithm 1.

To jointly learn the SiMT policy (i.e., aligned positions which determine latency) with translation, we weaken the attention of source words far away from the predicted aligned position in advance, thereby forcing the model to move the predicted aligned position to the source word that is most informative for translation (i.e., with the highest soft attention).

To this end, for $i^{th}$ target word, we introduce a Gaussian distribution $\bm{\mathcal{G}_{i}}$ centered on aligned position $p_{i}$ as the prior probability, calculated as:

where the Gaussian distribution is limited in first $g(i)$ source words. $\sigma_{i}$ is the variance used to control the attenuation degree of the prior probability as away from the aligned position. To prevent the prior probability of the furthest source word from being too small, we set $\sigma_{i}={p_{i}}/{2}$ according to the “two-sigma rule” Pukelsheim (1994). We will compare the performance of different settings of prior probability in Sec.6.1. Note that since the source position is discrete, we normalize the Gaussian distribution with $\mathcal{G}_{ij}/\sum_{k=1}^{g(i)}\mathcal{G}_{ik}$.

Given the prior probability $\mathcal{G}_{ij}$ and soft attention $\alpha_{ij}$ (calculated as Eq.(2)), which is considered as likelihood probability, we calculate the posterior probability $\widehat{\beta}_{ij}$ and normalize it as the final attention distribution $\beta_{ij}$:

Then, the context vector $c_{i}$ is calculated as:

When GMA is integrated into the Transformer with $L$ decoder layers and $H$ attention heads per layer, if multiple heads (totally $L\times H$ heads) independently predict their alignments, some outlier¹¹1Outlier heads mean that most of the heads are aligned in the front position, while some individual heads are aligned to the farther position, which requires the model to wait until the farthest aligned word is received, causing unnecessary latency. heads will cause unnecessary latency Ma et al. (2020); Zhang and Feng (2022a).

Therefore, to better adapt to multi-head attention and capture alignments, for each decoder layer, $H$ heads in GMA jointly predict the aligned source position and share it among $H$ heads, while the predicted alignments in each decoder layer still remain independent. Since the output position (determined by predicted alignments) in each layer may be different, the model starts translating after reaching the furthest one. We will compare the performance of different sharing settings in Sec.6.2.

We evaluate GMA on the following public datasets.

IWSLT15²²2nlp.stanford.edu/projects/nmt/ English $\!\rightarrow\!$ Vietnamese (En$\rightarrow$Vi) (133K pairs) Cettolo et al. (2015) We use TED tst2012 as validation set (1553 pairs) and TED tst2013 as test set (1268 pairs). Following the previous setting Raffel et al. (2017); Ma et al. (2020), we replace tokens that the frequency less than 5 by $\left\langle unk\right\rangle$, and the vocabulary sizes are 17K and 7.7K for English and Vietnamese respectively.

WMT15³³3www.statmt.org/wmt15/translation-task German $\!\rightarrow\!$ English (De$\rightarrow$En) (4.5M pairs) Following Ma et al. (2019), Arivazhagan et al. (2019) and Ma et al. (2020), we use newstest2013 as validation set (3000 pairs) and newstest2015 as test set (2169 pairs). BPE Sennrich et al. (2016) was applied with 32K merge operations and the vocabulary is shared across languages.

We conduct experiments on the following systems.

Offline Conventional Transformer Vaswani et al. (2017) model for full-sentence translation.

Wait-k Wait-k policy proposed by Ma et al. (2019), the most widely used fixed policy with strong performance and simple structure, which first waits for $k$ source words and then translates a target word and waits for a source word alternately.

MU A segmentation policy proposed by Zhang et al. (2020), which classify whether the current inputs is a complete meaning unit (MU), and then fed MU into the full-sentence MT model for translation until generating $\left\langle\mathrm{EOS}\right\rangle$. We compare our method with MU on De$\rightarrow$En(Big) since they report their results on De$\rightarrow$En with Transformer-Big.

MMA⁴⁴4github.com/pytorch/fairseq/tree/master/examples/simultaneous_translation Monotonic multi-head attention (MMA) proposed by Ma et al. (2020), the state-of-the-art adaptive policy. At each step, MMA predicts a Bernoulli variable to decide whether to start translating or wait for the next source token.

GMA Proposed method in Sec.3.

The implementations of all systems are adapted from Fairseq Library Ott et al. (2019) based on Transformer Vaswani et al. (2017). The setting is the same as Ma et al. (2020). For En$\rightarrow$Vi, we apply Transformer-small (6 layers, 4 heads per layer). For De$\rightarrow$En, we apply Transformer-Base (6 layers, 8 heads per layer) and Transformer-Big (6 layers, 16 heads per layer). We evaluate these systems with BLEU Papineni et al. (2002) for translation quality and Average Lagging (AL) Ma et al. (2019) for latency. Average lagging is currently the most widely used latency metric, which evaluates the number of words lagging behind the ideal policy. Given $g\left(i\right)$, AL is calculated as:

where $\left|\mathbf{x}\right|$ and $\left|\mathbf{y}\right|$ are the length of the source sentence and target sentence respectively.

We compared GMA with the Wait-k, MU and MMA, the current best representative of fixed policy, segmentation policy and adaptive policy respectively, and plot latency-quality curves in Figure 3, where GMA curve is drawn with various $\delta$ (in Eq.(6)), Wait-k curve is drawn with various lagging numbers $k$, MU curve is drawn with various classification thresholds of the meaning unit, MMA curve is drawn with various latency loss weights $\lambda$.

Compared with Wait-k, GMA has a significant improvement, since Wait-k ignores the alignments and thus the target word may be forced to be translated before receiving its aligned source word, which seriously affects the translation quality. Compared with MMA and MU, our method achieves better performance under most latency levels. Since MU first segments the source sentence based on the meaning unit, and then translates each segment with the full-sentence MT model, MU performed particularly well under high latency, but meanwhile it is difficult to extend to lower latency.

Compared with the SOTA adaptive policy MMA, GMA has stable performance and simpler training method. MMA introduces two additional loss functions to control the latency, and meanwhile applies the expectation training to train Bernoulli variables Ma et al. (2020). GMA successfully balances the translation quality and latency without any additional loss function. Owing to the proposed Gaussian prior probability centered on predicted alignments, the source words far away from the aligned position get less Gaussian prior, so that the model is forced to move the aligned position close to the most informative source word, thereby capturing the alignments and controlling the latency. With GMA, the translation quality and latency are integrated into a unified manner and jointly optimized without any additional loss function.

We conduct sufficient ablation studies on the method of modeling aligned position and the proposed prior probability in Table 1. For modeling aligned source position, predicting incremental step performs better. Since the complete length of the streaming inputs in SiMT is unknown, predicting absolute position easily exceeds the source length, resulting in unnecessary latency. In practice, the value range of incremental step is more regular than absolute position and thus easier to learn.

Among the prior probabilities of different distributions, the Gaussian distribution performs best. The Laplace distribution is more fat-tailed than the Gaussian distribution (i.e., more prior probability on the position that far away from the aligned position), resulting in learning a later aligned position and higher latency. The linear distribution performs worse since its attenuation with distance is smoother. When removing the prior probability, since the parameters to predict alignments do not get the back-propagation gradient, the model cannot learn the alignments at all, resulting in very low latency and poor translation quality. Focusing on the best performing Gaussian prior probability, when $\sigma_{i}=p_{i}/1$ (the attenuation degree is small), the predicted aligned position is much later and the latency is higher. when $\sigma_{i}=p_{i}/3$, the translation quality declines since the prior of some source words far away from the aligned position is too small. When the $\sigma_{i}$ is predicted, some small predicted $\sigma_{i}$ will make the prior probability of the distant source words almost 0, resulting in poor translation quality. In comparison, the Gaussian prior ($\sigma_{i}=p_{i}/2$) we proposed not only guarantees a certain attenuation degree, but also assigns the furthest source word some prior probability.

When integrated into multi-head attention, to reduce the overall latency of the model, GMA shares the predicted alignments among $H$ heads in each layer. Table 2 reports the performance of sharing alignments in different parts (all independent, among heads, among layers or among all).

‘All independent’ achieves the best translation quality and also brings a higher latency, as the overall latency of the model is determined by the farthest one among all predicted positions. Besides, ‘Share all’ gets the lowest latency but loses translation quality. Comparing ‘Share among heads’ and ‘Share among layers’, sharing among heads performs better, which is in line with the previous conclusion that there are obvious differences between the alignments in each decoder layer Garg et al. (2019); Voita et al. (2019); Li et al. (2019).

Unlike MMA predicting the READ/WRITE action, GMA directly predicts the incremental step between adjacent target outputs. To analyze the advantages of modeling the incremental step, we show the distribution of step size (i.e., the number of waiting source words between two adjacent outputs) in Figure 4, where we select the SiMT models with the similar latency (AL $\approx$ 4.5).

For Wait-k, the step size is blunt and there are only three cases, which occur before the start of translation (step size$=\!k$: first lag $k$ words), during translation (step size$=\!1$: wait and output one word alternately.), and after the end of the source inputs (step size$=\!0$: output the translation at one time). The proposed GMA and MMA have obvious differences in the step size distribution. The step size distribution of GMA is more even, only distributed between 0$\sim$4, which shows that each source segment is shorter and thus the translating process is more streaming. The step size of MMA is more widely distributed, most of which are 0 (consecutively output target words), and there is also a large proportion of step size greater than 5, which shows that MMA tends to consecutively wait for more source words and then output more target words, resulting in longer source segments. Therefore, although GMA and MMA have similar latency (AL), their performance in real applications is different, where the translation process of GMA is more streaming, while MMA is more segmented.

Furthermore, to more accurately evaluate the streaming degree of the translation process, we apply Consecutive Wait (CW) Gu et al. (2017) as the latency metric to evaluate the systems. Consecutive wait evaluates the number of source words waited between two target words, which reflects the streaming degree of SiMT. Given $g\left(i\right)$, CW is calculated as:

where $\mathbbm{1}_{g(i)-g(i-1)>0}=1$ counts the number of $g(i)-g(i-1)>0$. In other words, CW measures the average source segment length (the best case is 1 for word-by-word streaming translation and the worst case is $\left|\mathbf{x}\right|$ for full-sentence MT), where the smaller the CW, the shorter the average source segment, and the translation is more streaming.

As shown in Figure 5, GMA gets much smaller CW scores than MMA, where the average source segment length is about 2 words, which shows that GMA achieves more streaming translation than the previous adaptive policy.

Additionally, although the alignments between the two languages is not necessarily monotonic, we require the predicted incremental step $\Delta p_{i}\geq 0$ to guarantee the model performs READ/WRITE monotonically. This is due to two considerations. First, we don’t want to waste any useful source content, i.e., to avoid the model moving $p_{i}$ to the previous position and thereby ignoring some received source information caused by $\Delta p_{i}<0$. Second, we argue that monotonic alignments are more friendly to SiMT learning. We use fast-align⁵⁵5https://github.com/clab/fast_align Dyer et al. (2013) to generate the ground-truth aligned source position of $i^{th}$ target token, denoted as $A_{i}$, and then show the distribution of distances between adjacent alignments under monotonic and non-monotonic alignments in Figure 6, where ‘Non-Monotonic’ measures $A_{i}\!-\!A_{i-1}$ and ‘Monotonic’ measures $\mathrm{max}(A_{i}\!-\!\mathrm{max}_{j<i}A_{j},0)$. The distance distribution between adjacent alignments under monotonic alignment is more concentrated, between 0 and 4, which is easier for the model to learn incremental step. Actually, the incremental steps predicted by GMA almost distribute between 0 and 4 (see Figure 4), which shows that GMA successfully learns the relative distance between monotonic alignments.

To evaluate the quality of the aligned source position $p_{i}$ predicted by GMA, we measure the alignment accuracy on the RWTH De$\rightarrow$En alignment dataset ⁶⁶6https://www-i6.informatik.rwth-aachen.de/goldAlignment/, whose reference alignment was manually annotated by experts Liu et al. (2016); Zhang and Feng (2021b). As shown in Table 3, we sample one decoder layer to calculate the alignments error rate (AER) Vilar et al. (2006), and meanwhile count how many ground-truth aligned source words are located before the output position $g(i)$ (i.e., translate after receiving the aligned source word).

GMA achieves good alignment accuracy, especially at low latency, since the model is required to output immediately after receiving the aligned source words ($\delta$ in Eq.(6) is small). More importantly, most of the ground-truth alignments is within $g(i)$, showing that GMA guarantees that in most cases, the model starts translating a target word after receiving its aligned source words, which is beneficial to translation quality.

To study how GMA learns to balance translation quality and latency without any additional loss function during training, we draw a learning curve for translation quality and latency in Figure 7.

Initially, the high latency indicates that the model first moves the predicted aligned position to a further position, to learn the translation by seeing more source words. Then, as the number of training steps increases, the translation quality improves and the latency gradually decreases, which shows that for better translation, the model moves the predicted aligned source position to a more appropriate position due to the introduced Gaussian prior probability. Overall, for better translation, GMA constantly adjusts the predicted aligned source position to a suitable position and thereby controls the latency, which is completely different from the previous method of introducing the additional latency loss to constrain the latency with translation Ma et al. (2020); Miao et al. (2021).

We explore the characteristics of GMA by visualizing the attention distributions in Figure 8. We show two cases with the alignments of different difficulty levels, where the reverse orders in alignments are considered as a major challenge for SiMT Ma et al. (2019); Zhang and Feng (2021c).

For more monotonic alignments, GMA can predict the aligned source position well and output the target word after receiving the aligned word. Meanwhile, due to the characteristics of Gaussian distribution, GMA can also avoid focusing too much on source words in the front position, and strengthen the attention on the newly received source words, which is proved to be beneficial to SiMT performance Elbayad et al. (2020b); Zhang and Feng (2022b). For more complex alignments, the aligned position predicted by GMA is close to the ground-truth alignments, so that GMA starts translating after receiving most of aligned words. Besides, GMA learns some implicit prediction ability, e.g., before receiving “worden sein”, GMA generates the correct translation “have been” based on the context. We consider this is because the predicted alignments during training are monotonic due to the incremental step, where modeling monotonic alignments forces the model to learn the correct translation from the incomplete source and previous outputs Ma et al. (2019).