Sharing Attention Weights for Fast Transformer

1 Introduction

In recent years, neural models have led to great improvements in machine translation (MT). Approaches of this kind make it possible to learn good mappings between sequences via deep networks and attention mechanisms [Sutskever et al., 2014; Bahdanau et al., 2015; Luong et al., 2015]. Recent work has explored an architecture that just consists of stacked attentive and feed-forward networks (call it Transformer) [Vaswani et al., 2017]. It makes use of multi-layer dot-product attention to capture the dependency among language units. Beyond this, training this kind of model is fast because we can parallelize computation over all positions of the sequence on modern graphics processing units (GPUs). These properties make Transformer popular in recent MT evaluations and industrial deployments.

However, standard implementations of Transformer are prone to slow inference though fast in training. At test time, the system produces one target word each time until an end symbol is reached. This process is auto-regressive and slow because we have to run dot-product attention for each position rather than batching the computation of the sequence.

The situation is even worse if 6 or more attention layers are stacked and the attention model occupies the inference time. To address this issue, efficient networks have been investigated. For example, one can replace dot-product attention with additive attention and use average attention models instead [Zhang et al., 2018], or explore non-autoregressive decoders that benefit from the trick of batched matrix operations over the entire sequence [Gu et al., 2018]. But these methods either lose the explicit model of word dependencies, or require complicated networks that are hard to train.

In this work, we observe that the attention model shares a similar distribution among layers in weighting different positions of the sequence. This experience lead us to study the issue in another line of research, in which we reduce redundant computation and re-use some of the hidden states in the attention network. We propose a method to share attention weights in adjacent layers (call it shared attention network, or SAN for short). It leads to a model that shares attention computation in the stacked layers vertically. In addition to the new architecture, we develop a joint method to learn sharing policies and MT models simultaneously. As another "bonus", SAN reduces the memory footprint because some hidden states are kept in the same piece of memory.

SAN is simple and can be implemented in a few hours by anyone with an existing kit of Transformer. Also, it is orthogonal to previous methods and is straightforwardly applicable to the variants of Transformer. We test our approach in a stateof-the-art system where an attention cache is already in use for speed-up. Experimental results on ten WMT and NIST OpenMT tasks show an average of 1.3X speed-up with almost no decrease in BLEU. More interestingly, it obtains a bigger speed-up (1.8X) when working with the AAN model. The best result is 16 times faster than the baseline where no cache is adopted.

2 The Transformer System

The Transformer system follows the popular encoder-decoder paradigm. On the encoder side, there are a number of identical stacked layers. Each of them is composed of a selfattention sub-layer and a feed-forward sub-layer. In Transformer, the attention model is scaled dot-product attention. Let l be the length of the source sequence. The input of the attention sub-layer is a tuple of (Q, K, V ), where Q ∈ Rl×dk , K ∈ Rl×dk and V ∈ Rl×dv are the matrices of queries, keys, and values packed over the sequence. In self-attention, we first compute the dot-product of queries and keys, followed by the rescaling and softmax operations.

S is an l × l matrix, where entry (i, j) represents the strength of connecting position i with position j. Note that S is essentially a weight (or scalar) matrix where every column represents a distribution. The output of self-attention is simply defined as the weighted sum of values:

Here Q, K and V are generated from the same source with a linear transformation. The self-attention result is then fed into a fully connected feed-forward network (FFN).

The decoder shares a similar structure with the encoder. Apart from the self-attention sub-layer, an encoder-decoder attention sub-layer is introduced to model the correspondence between source positions and target positions. Basically, the encoder-decoder attention has the same form as Eqs. (1) and (2), where the queries come from the output of the previous layer, and the keys and values come from the output of the encoder. See Figure 1 for an illustration of the attention model used in Transformer.

Note that the matrix multiplications in Eqs. (1) and (2) are time consuming. It is a bigger problem for inference because Eqs. (1) and (2) repeat for each position until we finish the generation.

3 Shared Attention Networks

In this work we speed up the decoder-side attention because the decoder is the heaviest component in Transformer.

3.1 AttentionWeights

Self-attention is essentially a procedure that fuses the input values to form a new value at each position. Let S[i] be column i of weight matrix S. For position i , we first compute S[i] to weight all positions (as in Eq. (1)), and then compute the weighted sum of values by S[i] (as in Eq. (2)). In column vector S[i], element Si,j indicates the contribution that we fuse the value at position j to position i. Intuitively, the attention weight S[i] should not be volatile in different levels of language representation because the correlations between positions somehow reflect the dependency of language units. For example, for an English sentence, the subject and the verb correlate well no matter how many layers we make on top of the input sequence. On the other hand, the subject and the adverbial may not have a big impact to each other in all stacked layers.

To verify this, we compute the Jensen-Shannon (JS) divergence to measure how the attention weight distribution of a layer is different from another [Lin, 1991]. We choose the JS divergence because it is symmetric and bounded. For multihead attention, we regard different heads as separate channels. We compute the JS score for each individual head and then average them for final output. Figure 2 shows that the system generates similar weights over layers. For self-attention, layers 2-6 almost enjoy the same weight distribution. For encoder-decoder attention, we observe a larger variance but good similarities still exist among two or three adjacent layers (see entries around the diagonal of Figure 2(b)). All these show the possibility of removing redundant computation in Transformer.

3.2 The Model

An obvious next step is to develop a faster attention model that makes efficient re-use of the states in Eqs. (1) and (2),

· · ·

· · ·

instead of computing everything on the fly. In this work we present a shared attention network (SAN) to share weight matrix S for adjacent layers. The idea is that we just compute the weight matrix once and reuse it for upper-level layers. Here we describe SAN for both the self-attention and encoderdecoder attention models.

• SAN Self-Attention. We define the self-attention weight matrix in layer m as:

where s(·, ·) is the function described in Eq. (1), Qm and Km are the inputs, and Sm is the attention weight for the output. In SAN, we can share Sm with the layers above m, like this

where π indicates how many layers share the same attention weights. For example, in a 6-layer decoder, we can share the self-attention weights for every two layers (π = 2), or share the weights for the first two layers (π1 = 2) and let the remaining 4 layers use another weights (π2 = 4). We discuss the sharing strategy in the later part of the section.

• SAN Encoder-Decoder Attention. For encoderdecoder attention, we do the same thing as in selfattention, but with a trick for further speed-up. In encoder-decoder attention, keys and values are from the output of the encoder, i.e., K and V have already been shared among layers on the decoder side. In response, we can share A = S · V for encoder-decoder attention layers. This can be described as

where Am is the attention output of layer m, V is the context representation generated by the encoder. See Figure 3 for a comparison of the standard attention model and SAN.

In addition to system speed-up, SAN also reduces the memory footprint. In SAN, we just need one data copy of weight matrix for a layer block, rather than allocating memory space for every layer. Moreover, the linear transformation of the input (i.e., Q and K) can be discarded when the attention weights come from another layer. It reduces both the number of model parameters and the memory footprint in inference.

Another note on SAN. SAN is a process that simplifies the model and re-uses hidden states in the network. It is doing something similar to systems that share model parameters in different levels of the network [Wu et al., 2016; Yang et al., 2018; Luong et al., 2016]. Such methods have been proven to improve the robustness of neural models on many natural language processing tasks. Sharing parameters and/or hidden states can reduce the model complexity. Previous work has pointed out that MT systems cannot benefit a lot from very deep and complex networks [Vaswani et al., 2017; Britz et al., 2017]. SAN might alleviate this problem and makes it easier to train neural MT models. For example, in our experiments, we see that SAN can improve translation quality in some cases in addition to considerable speed-ups.

3.3 Learning to Share

The remaining issue is how to decide which layers can be shared. A simple way is to use the same setting of π for the

· · ·

entire layer stack, and tune it on a development set. For example, we can try to share weights on layer blocks consisting of two layers, or three layers, or all layers (π = 2, or 3 , ...), and use the tuned π on test data.

But a uniform sharing strategy might not be optimal because we need to control the degree of sharing in difference levels of the network. For example, for the case in Figure 2(a), a good choice is to share weights for layers 2-6 and leave layer 1 as it is. Here we present a method that learns the sharing strategy in a dynamic way. To do this, we choose ln(2)- JS divergence as the measure of the similarity between weights of layer i and layer j (denoted as µ(i, j)). Given a layer block ranging from layer m to layer n (denoted as b(m, n)), the similarity over the block is defined as

where n-m+1 is the size of the block, and δ(i, j) is the Kronecker delta function. sim(m, n) measures how the weight of a layer is similar to that of another layer in block b(m, n). We can do sharing when sim(m, n) ≥ θ where θ is the parameter that controls how often a layer is shared.

We begin with layer 1 and search for the biggest block that satisfies the criterion. This process repeats until all the layers are considered, resulting in N layer blocks. For simplicity, we use {π1, ...,πN} (or {πi}) to represent the blocks in a bottom-up manner (call it sharing policy), where πi is the size of block i. Obviously, for an M-layer stack we have ∑Ni=1 πi = M.

Once we have a sharing policy, we need to re-train the MT model. It in turn leads to new attention weights and possibly a better policy. A desirable way is to continue learning until the model converges. To this end, we present a joint learning method that trains MT models and sharing policies simultaneously (Figure 4). In the method, MT training and policy learning loops for iterations, and the result of the final round is returned when there is no new update of the model.

4 Experiments

We experimented with our approach on WMT and NIST translation tasks.

4.1 Experimental Setup

The bilingual and evaluation data came from three sources

• WMT14 (En-De). We used all bilingual data provided within the WMT14 English-German task. We chose newstest 2013 as the tuning data, and newstest 2014 as the test data.
• WMT17 (En-De, De-En, En-Fi, Fi-En, En-Lv, Lv-En, En-Ru and Ru-En). We followed the standard data setting of the bidirectional translation tasks of GermanEnglish, Finnish-English, Latvian-English, and RussianEnglish. For tuning, we concatenated the data of newstest 2014-2016. For test, we chose newstest 2017.
• NIST12 (Zh-En). We also used parts of the bitext of NIST OpenMT12 to train a Chinese-English system 1 . The tuning and test sets were MT06 and MT08.

For Chinese, all sentences were word segmented by the segmentation system in the NiuTrans toolkit [Xiao et al., 2012]. For other languages, we ran the official script of WMT for tokenization. All sentences of more than 50 words were removed for the NIST Zh-En task, and sentences of more than 80 words were removed for the WMT tasks. For all these tasks, sentences were encoded using byte-pair encoding, where we used a shared source target vocabulary of 32K tokens. See Table 1 for statistics of the data.

We used standard implementation of Transformer. Early versions of its inference system simply compute the attention output for target positions individually. This way is straightforward but with a double counting problem. For a stronger baseline, we chose the system with an attention cache that kept the attention output of previous positions in cache and re-used it in following steps.

The Transformer system used in our experiments consisted of a 6-layer encoder and a 6-layer decoder. By default, we set dk = dv = 512 and used 2,048 hidden units in the FFN sub-layers. We used multi-head attention (8 heads) because it was shown to be effective for state-of-the-art performance [Vaswani et al., 2017]. Dropout (rate = 0.1) and label smoothing (ϵls = 0.1) methods were adopted for regularization and stabilizing the training [Szegedy et al., 2016]. We trained the model using Adam with β1 = 0.9, β2 = 0.98, and ϵ = 10-9 [Kingma and Ba, 2015]. The learning rate was scheduled as described in [Vaswani et al., 2017]: lr = d-0.5 ·min(t-0.5 , t·4k-1.5 ), where t is the step number. All models were trained for 100k steps with a mini-batch of 4,096 tokens on machines with 8 Nvidia 1080Ti GPUs except En-Fi (60k steps), Fi-En (60k steps) and Zh-En (24k steps). Every model was ensembled from the 5 latest checkpoints in training. For inference, both beam search and batch decoding methods were used (beam size = 4 and batch size = 16).

1 LDC2000T46, LDC2000T47, LDC2000T50, LDC2003E14, LDC2005T10, LDC2002E18, LDC2007T09 and LDC2004T08

For our approach, we applied SAN to self-attention and encoder-decoder sub-layers on the decoder side. We learned sharing policies as in Figure 4. θ was tuned on the tuning data, which resulted in an optimal range of [0.3, 0.4] for selfattention and [0.4, 0.5] for encoder-decoder attention.

4.2 Results

We report the translation quality (in BLEU[%]) and speed (in token/sec) on all ten of the tasks (Table 2). We see, first of all, that SAN significantly improves the translation speed for all these languages. The average speed-up is 1.3X. Also, there is a very modest BLEU decrease, but not significant. These results indicate that SAN is robust and can improve a strong baseline on a wide range of translation tasks. Another interesting finding here is that the speed improvement on En-Ru, Ru-En and Zh-En is not as large as that on other language pairs. This is because we share fewer layers (i.e., larger θ) on these tasks to preserve good BLEU scores. Note that Russian and Chinese are very difficult languages for translation, and we need a complicated network to model the structure divergence. Less sharing is preferred to keep the expressive power for these languages.

To modulate the degree of sharing, we study the system behavior under different settings of θ (Table 3) . For comparison, we report the result of uniform {πi}. Due to the limited space, we present the result on the WMT14 En-De task in the following sensitivity analysis. For uniform sharing, {πi} is set to {6} for self-attention and {3,3} for encoder-decoder attention. This results in a promising speed-up (see entry of uniform {πi}). When we switch to joint learning of MT models and sharing policies, we obtain further improvements in both BLEU and speed. More interestingly, we see that the system prefers a smaller θ for self-attention than the encoder-decoder counterpart. This is reasonable because the encoder-decoder attention captures the correspondence of two languages and needs more states in modeling than a single language. On the other hand, the BLEU improvements indicate that the MT system can benefit from simplified models. It gives a direction that we explore simpler models for better training of neural MT systems.

In encoder-decoder attention, we share the context V generated by the encoder for further speed-ups (see Figure 3(c)). It is therefore worth a study on how much this method can accelerate the system. Table 4 shows that sharing the context contributes half of the speed improvement. This agrees with our design that weight sharing is more beneficial to the decoder because attention is heavier on the decoder side. Another interesting question is whether SAN can improve the system on the encoder side. To seek an answer, we apply SAN to the encoder-side self-attention sub-layers and see small speed improvements (Table 5). This result confirms the previous report that the decoder occupies the inference time and the encoder is light [Zhang et al., 2018].

Also, we plot the translation speed as functions of beam size and BLEU score. Figure 5 shows a consistent improvement under different beam settings. Moreover, SAN benefits more from larger beam sizes. For example, the speed-up of beam = 20 is larger than that of beam = 4 (1.48x vs. 1.40x). The Speed-vs-BLEU curves indicate a good ability of SAN in trading off between translation quality and speed.

In addition, we empirically compare SAN with other variants of the attention model, including AAN [Zhang et al., 2018] and the model with no cache. Table 6 shows the attention cache plays an important role in fast inference. It leads to an 8-fold speed-up on top of the implementation where no cache is used. Also, SAN obtains a bigger speed improvement than AAN. This might be because AAN is used for selfattention only, while SAN is applicable to both self-attention and encoder-decoder attention. Finally, we combine AAN and SAN in a new system where AAN is applied to self-attention and SAN is applied to encoder-decoder attention. It yields the best result which is 1.8 times faster than the baseline, and almost 16 times faster than the systemwithout cache.

In training, we observe that systems tend to learn similar attention weights. Figure 6 plots the JS divergence between layer 4 and layers 5-6 at different training steps. The JS divergence curves go down significantly as the training proceeds. Adjacent layers show more similar weight distributions than distant layers. The fast convergence in JS divergence can speed up the iterative training. For example, for each training epoch (Figure 4), one can train the model for a shorter time, as the JS divergence among layers converges quickly. Thus, the system can find the optimal sharing policy more efficiently. In addition, we find that the training likelihood of SAN is higher than that of the baseline, but not significant.

5 RelatedWork

It has been observed that attention models are critical for state-of-the-art results on many MT tasks [Bahdanau et al., 2015; Wu et al., 2016; Vaswani et al., 2017]. Several research groups have investigated attentive methods for different architectures of neural MT. The earliest is [Luong et al., 2015]. They introduced an additive attention model into MT systems based on recurrent neural networks (RNNs). More recently, multi-layer attention was successfully applied to convolutional neural MT systems [Gehring et al., 2017] and Transformer systems [Vaswani et al., 2017]. In particular, Transformer is popular due to its scalability on large-scale training and the good design of the architecture for implementation.

It is well-known that Transformer suffers from a high inference cost which makes it slower than the RNN-based counterpart. This partially due to the auto-regressive property of decoding, and partially due to the heavy use of dot-product attention where the expensive matrix multiplication is frequently used. Researchers have begun to explore solutions. For example, Gu et al. [2018] designed a non-autoregressive inference method for a Transformer-like system, which generated the entire target sequence at one time. This model is fast but is not easy to train.

In another line of research, Zhang et al. [2018] proposed the average attention network (AAN) and applied it to selfattention sub-layers on the decoder side with no cache. In this work, we study the issue on a strong baseline that has already used an attention cache for a reasonable inference speed. Also, our approach is straightforwardly applicable to systems with multi-layer attention. We improve both the selfattention and encoder-decoder attention components, which has not been studied in previous work.

In neural MT, fast inference methods have been investigated for years. These include vocabulary selection [L'Hostis et al., 2016; Sankaran et al., 2017], knowledge distillation [Hinton et al., 2015; Kim and Rush, 2016], low-precision computation [Micikevicius et al., 2018; Quinn and Ballesteros, 2018], recurrent stacked networks [Dabre and Fujita, 2019] and etc. Our method is orthogonal to them. Previous studies focus more on the reduction of model size and robust training, rather than fast inference. Here we study the issue in the context of speeding up attentive MT and confirm the effectiveness of this kind of models.

6 Conclusions

We have presented a shared attention network (SAN) for fast inference of Transformer. It shares attention weights among layers for both self-attention and encoder-decoder attention in a vertical manner. The policy of sharing can be jointly learned with the MT model, rather than being determined heuristically. Moreover, SAN reduces the memory footprint. Experiments on ten MT tasks show that SAN yields a speed-up of 1.3X over a strong baseline that has already used an attention cache. More interestingly, it is observed that the combination of SAN and AAN obtains a larger speed improvement. The system is 16X faster than the baseline with no cache.

Acknowledgments

This work was supported in part by the National Science Foundation of China (Nos. 61732005, 61876035 and 61432013) and the Fundamental Research Funds for the Central Universities (No. N181602013).