1 Introduction↩︎

Transformer is an attention-based encoder-decoder model [1]. It has shown promising results in machine translation tasks recently [2]–[5]. Nonetheless, Transformer suffers from the inefficiency issue at inference. This problem is attributed to the Transformer decoder for two reasons: 1) the decoder is deep [6]. It consists of multiple layers and each layer contains three sub-layers, including two attentions and a feed-forward network; 2) the attention has a high (quadratic time) complexity [7], as it needs to compute the correlation between any two input words.

Previous work has focused on improving the complexity of the attention in the decoder to accelerate the inference. For example, Aan uses the averaging operation to avoid computing the correlation between input words [7]. San share the attention results among layers [8]. On the other hand, we learn that vanilla attention runs faster in training than in inference thanks to its parallelism. This offers a new direction: a higher degree of parallelism could speed up the inference. The most representative work of this type is the non-autoregressive approach [9]. Its decoder predicts all words in parallel, but fails to model the word dependencies. Despite of their successes, all these systems still have a deep decoder.

In this work, we propose to parallelize the sub-layers to obtain a shallow autoregressive decoder. This way does not suffer from the poor result of directly reducing depths and avoids the limitation of non-autoregressive approaches. We prove that the two attention sub-layers in a decoder layer could be parallelized if we assume their inputs are close to each other. This assumption holds and thereby we compress these two attentions into one. Furthermore, we show that the remaining feed-forward network could also be merged into the attention due to their linearity. To the end, we propose Compressed Attention Network (Can for short). The decoder layer of Can possesses a single attention sub-layer that does the previous three sub-layers’ jobs in parallel. As another “bonus”, Can is simple and easy to be implemented.

In addition, [6] empirically discover that existing systems are not well balancing the encoder and decoder depths. Based on their work, we build a system with a deep encoder and a shallow decoder, which is 2\(\times\) faster than the widely used standard baseline without loss in performance. It requires neither the architecture modification nor adding extra parameters. This system serves as a stronger baseline for a more convincing comparison.

We evaluate Can and the stronger baseline in 14 machine translation tasks, including WMT14 English\(\leftrightarrow\) (En\(\leftrightarrow\)) and WMT17 English\(\leftrightarrow\) (En\(\leftrightarrow\)). The experiments show that Can is up to 2.82\(\times\) faster than the standard baseline with almost no loss in performance. Even comparing to our stronger baseline, Can still has a 1.42\(\times\) speed-up, while other acceleration techniques such as San and Aan are 1.12\(\sim\)​1.16\(\times\) in the same case.

To summarize, our contributions are as follows:

• We propose Can, a novel architecture that accelerates Transformer by compressing its sub-layers for a higher degree of parallelism. Can is easy to be implemented.

• Our work is based on a stronger baseline, which is 2\(\times\) faster than the widely used standard baseline.

• The extensive experiments on 14 WMT machine translation tasks show that Can is 1.42\(\times\) faster than the stronger baseline and 2.82\(\times\) for the standard baseline. Can also outperforms other approaches such as San and Aan.

2 Background: Transformer↩︎

Transformer is one of the state-of-the-art neural models in machine translation. It consists of a \(N\)-layer encoder and a \(N\)-layer decoder, where \(N=6\) in most cases. The encoder maps the source sentence to a sequence of continuous representations and the decoder maps these representations to the target sentence. All layers in the encoder or decoder are identical to each other.

The layer in the decoder consists of three sub-layers, including the self-attention, the cross-attention and the feed-forward network (FFN). The self-attention takes the output \(X\) of the previous sub-layer as its input and produces a tensor with the same size as its output. It computes the attention distribution \(A_x\) and then averages \(X\) by \(A_x\). We denote the self-attention as \(Y_x=\mathrm{Self}(X)\), where \(X\in\mathbb{R}^{t\times d}\), \(t\) is the target sentence length and \(d\) is the dimension of the hidden representation: \[\begin{align} A_x&=\mathrm{SoftMax}(\frac{XW_{q1}W_{k1}^TX^T}{\sqrt{d}})\tag{1}\\ Y_x&=A_xXW_{v1}\tag{2} \end{align}\] where \(W_{q1},W_{k1},W_{v1}\in\mathbb{R}^{d\times d}\).

The cross-attention is similar to the self-attention, except that it takes the encoder output \(H\) as an additional input. We denote the cross-attention as \(Y_h=\mathrm{Cross}(X,H)\), where \(H\in\mathbb{R}^{s\times d}\), \(s\) is the source sentence length: \[\begin{align} A_h&=\mathrm{SoftMax}(\frac{XW_{q2}W_{k2}^TH^T}{\sqrt{d}})\tag{3}\\ Y_h&=A_hHW_{v2}\tag{4} \end{align}\] where \(W_{q2},W_{k2},W_{v2}\in\mathbb{R}^{d\times d}\).

The FFN applies non-linear transformation to its input \(X\). We denote the FFN as \(Y_f=\mathrm{FFN}(X)\): \[Y_f=\mathrm{ReLU}(XW_1+b_1)W_2+b_2\label{eqn:ffn}\tag{5}\] where \(W_1\in\mathbb{R}^{d\times 4d}\), \(b_1\in\mathbb{R}^{4d}\), \(W_2\in\mathbb{R}^{4d\times d}\) and \(b_2\in\mathbb{R}^{d}\).

All sub-layers are coupled with the residual connection [10], i.e., \(Y=f(X)+X\) where \(f\) could be any sub-layer. Their inputs are also preprocessed by the layer normalization first [11]. Fig. [fig:transformer] shows the architecture of Transformer decoder. For more details, we refer the reader to [1].

3 Compressed Attention Network↩︎

3.1 Compressing Self-Attention and Cross-Attention↩︎

As suggested by [12], the output of one layer in the residual network can be decomposed into the sum of all outputs from previous layers. For the adjacent self-attention and cross-attention, we can write their final output as \(Y=X+\mathrm{Self}(X)+\mathrm{Cross}(X',H)\), where \(X\) is the input of self-attention and \(X'=X+\mathrm{Self}(X)\) is the input of cross-attention. If \(X\) and \(X'\) are identical, we are able to accelerate the computation of \(Y\) by parallelizing these two attentions, as \(X'\) do not need to wait \(\mathrm{Self}(X)\) to finish.

Previous work [13] has shown that inputs of adjacent layers are similar. This implies that \(X\) and \(X'\) are close and the parallelization is possible. We empirically verify this in the left part of Fig. 2 by examining the cosine similarity between inputs of every self-attention and cross-attention pairs. It shows that \(X\) and \(X'\) are indeed close to each other (a high similarity \(>0.9\) for the diagonal entries). Therefore we could assume \(X\) and \(X'\) are identical (we omit the layer normalization for simplicity): \[Y=X+\mathrm{Self}(X)+\mathrm{Cross}(X,H)\label{eqn:self-cross-res}\tag{6}\]

By observing that Eq. 2 and Eq. 4 are essentially matrix multiplications, we could rewrite \(\mathrm{Self}(X)+\mathrm{Cross}(X,H)\) as a single matrix multiplication: \[\begin{align} A&=\left[A_x^T,A_h^T\right]^T\\ \mathrm{Self}(X)+\mathrm{Cross}(X,H)&=A\left[XW_{v1},HW_{v2}\right]\label{eqn:self-cross-sum} \end{align}\tag{7}\] \([\cdot]\) is the concatenation operation along the first dimension.

[8] shows that some attention distributions \(A_x\) and \(A_h\) are duplicate. This means that there exists a certain redundancy in \(\{W_{q1},W_{k1}\}\) and \(\{W_{q2},W_{k2}\}\). Thus we could safely share \(W_{q1}\) and \(W_{q2}\) to parallelize the computation of the attention distribution \(A\): \[\begin{align} \bar{A}&=\left(XW_q\left[XW_{k1},HW_{k2}\right]^T\right)/\sqrt{d}\\ A&=\left[\mathrm{SoftMax}(\bar{A}^T_{\cdot,1\ldots t}),\mathrm{SoftMax}(\bar{A}^T_{\cdot,t+1\ldots t+s})\right]^T \end{align}\]

However, \(A\) consists of two SoftMax distributions and is used in Eq. 7 without normalization. The output variance is then doubled and leads to poor optimization [14]. It is advised to divide \(A\) by \(\sqrt{2}\) to preserve the variance. This way resembles a single distribution. So we use one SoftMax instead and this works well: \[A=\mathrm{SoftMax}(\frac{XW_q\left[XW_{k1},HW_{k2}\right]^T}{\sqrt{d}})\label{eqn:self-cross-weight}\tag{8}\]

Now, we can compute \(Y\) in Eq. 6 efficiently by using Eq. 8 as well as Eq. 7 to compute \(\mathrm{Self}(X)+\mathrm{Cross}(X,H)\).

3.2 Compressing Attention and FFN↩︎

It is natural to consider to merge the attention and FFN with the same approach for further speed-up. As suggested by the right part of Fig. 2, the similarities between inputs of the adjacent cross-attention and FFN are low (dark diagonal entries). This implies that it is not ideal to make the identical input assumption to parallelize the cross-attention and FFN.

Here we provide another solution. Given that attention is merely a weighted sum and FFN performs a linear projection first, we can merge them by exploiting the linearity. This way not only parallelizes the computation of attention and FFN but also removes redundant matrix multiplications.

We substitute \(X\) in Eq. 5 by \(Y\) in Eq. 6 : \[Y_f=\mathrm{ReLU}(XW_1+A\left[XW_{v1},HW_{v2}\right]W_1+b_1)W_2+b_2\]

We can combine \(W_1\) with \(W_{v1}\) as well as \(W_{v2}\) into \(\widetilde{W}_{v1},\widetilde{W}_{v2}\in\mathbb{R}^{d\times 4d}\), as these matrices are learnable and matrix multiplied together: \[Y_f=\mathrm{ReLU}(XW_1+A\left[X\widetilde{W}_{v1},H\widetilde{W}_{v2}\right]+b_1)W_2+b_2\label{eqn:self-cross-ffn}\tag{9}\] Furthermore, \(XW_1\) can be computed in parallel with other transformations such as \(XW_q\).

This eventually gives us an more efficient decoder layer architecture, named Compressed-Attention. The whole computation process is shown in Fig. 3: it first computes the attention distribution \(A\) by Eq. 8 , then performs the attention operation via Eq. 9 , and produces \(Y_f\) as the final result. The proposed Compressed Attention Network (Can) stacks compressed-attentions to form its decoder. Fig. 1 shows the difference between Transformer and Can.

3.3 Balancing Encoder and Decoder Depths↩︎

Based on the findings of [6], we learn that a shallow decoder could offer a great speed gain, while a deep encoder could make up of the loss of a shallow decoder without adding a heavy computation overhead. Since their work is based on knowledge distillation [15], here we re-examine this idea under the standard training setting (without knowledge distillation).

Fig. 4 shows the performance and speed if we gradually reduce the decoder depth while adding more encoder layers. We see that although the overall number of parameters remains the same, the baseline can be 2\(\times\) faster without losing any performance (12/2 vs. 6/6). This justifies the previous idea. We thereby choose a stronger baseline with a 12-layer encoder and a 2-layer decoder for a more convincing comparison. This setting is also applied to Can.

4 Experiments↩︎

4.1 Experimental Setup↩︎

4.1.1 Datasets↩︎

We evaluate our methods on 14 machine translation tasks (7 datasets \(\times\) 2 translation directions each), including WMT14 En\(\leftrightarrow\) and WMT17 En\(\leftrightarrow\).

WMT14 En\(\leftrightarrow\) datasets are tokenized by a script from Moses³. We apply BPE [16] with 32K merge operations to segment words into subword units. Sentences with more than 250 subword units are removed. The first two rows of Table [tab:data] are the detailed statistics of these two datasets. For En-De, we share the source and target vocabularies. We choose newstest-2013 as the validation set and newstest-2014 as the test set. For En-Fr, we validate the system on the combination of newstest-2012 and newstest-2013, and test it on newstest-2014.

All WMT17 datasets are the official preprocessed version from WMT17 website⁴. BPE with 32K merge operations is similarly applied to these datasets. We use the concatenation of all available preprocessed validation sets in WMT17 datasets as our validation set:

• En\(\leftrightarrow\)De. We use the concatenation of newstest2014, newstest2015 and newstest2016 as the validation set.

• En\(\leftrightarrow\)Fi. We use the concatenation of newstest2015, newsdev2015, newstest2016 and newstestB2016 as the validation set.

• En\(\leftrightarrow\)Lv. We use newsdev2016 as the validation set.

• En\(\leftrightarrow\)Ru. We use the concatenation of newstest2014, newstest2015 and newstest2016 as the validation set.

• En\(\leftrightarrow\)Cs. We use the concatenation of newstest2014, newstest2015 and newstest2016 as the validation set.

We use newstest2017 as the test set for all WMT17 datasets. Detailed statistics of these datasets are shown in Table [tab:data]. For all 14 translation tasks, we report case-sensitive tokenized BLEU scores⁵.

4.1.2 Model Setup↩︎

Our baseline system is based on the open-source implementation of the Transformer model presented in [17]. For all machine translation tasks, the standard Transformer baseline (Baseline) consists of a 6-layer encoder and a 6-layer decoder. The embedding size is set to 512. The number of attention heads is 8. The FFN hidden size equals to 4\(\times\) embedding size. Dropout with the value of 0.1 is used for regularization. We adopt the inverse square root learning rate schedule with 8,000 warmup steps and \(0.0007\) learning rate. We stop training until the model stops improving on the validation set. All systems are trained on 8 NVIDIA TITIAN V GPUs with mixed-precision training [18] and a batch size of 4,096 tokens per GPU. We average model parameters in the last 5 epochs for better performance. At test time, the model is decoded with a beam of width 4 and half-precision. For an accurate speed comparison, we decode with a batch size of 1 to avoid paddings. The stronger balanced baseline (Balanced) shares the setting with this standard baseline, except that its encoder depth is 12 and decoder depth is 2.

We compare Can and other model acceleration approaches with our baselines. We choose Sharing Attention Network (San) [8] and Average Attention Network (Aan) [7] for comparison, as they have been proven to be effective in various machine translation tasks [19]. All hyper-parameters of Can, San and Aan are identical to the balanced baseline system. Results are the average of 3 runs.

4.2 Results↩︎

Table [tab:wmt14] shows the results of various systems on WMT14 En\(\leftrightarrow\). Our balanced baseline has nearly the same performance as the standard baseline, but its speed is 2\(\times\) faster on average. A similar phenomenon is also observed from WMT17 experiments in Table [tab:wmt17]. This observation indicates that existing systems do not well balance the encoder and decoder depths. We also report the performance of Aan, San and the proposed Can. All three approaches have similar BLEU scores and slightly underperform the balanced baseline. Can is more stable than the others, as its maximum \(\Delta_{\mathrm{BLEU}}\) is -0.39, while San is -0.67 and Aan is -0.61. For speeds of these systems, San and Aan have a similar level of acceleration (1\(\sim\)​16%) over the balanced baseline. Can, on the other hand, provides a higher level of acceleration (27\(\sim\)​42%). Interestingly, we find that the acceleration is more obvious in De-En than in others, e.g., 42% in De-En and 27% in En-Fr for Can. We find that the length ratio between the translation and the source sentence in De-En is higher than others, e.g., 1.0 for De-En and 0.981 for En-Fr. In this case the decoder tends to predict more words and consumes more time in De-En, and thus acceleration approaches that work on the decoder are more effective. In addition, though not reported in Table [tab:wmt14], we find that applying Can on the standard baseline hurt the performance less (-0.09 BLEU points on average) than on the balanced baseline (-0.25 BLEU points on average).

More experimental results to justify the effectiveness of Can are presented in Table [tab:wmt17]. We evaluate the balanced baseline as well as Can on five WMT17 language pairs. The results again show that the balanced baseline is indeed a strong baseline with BLEU scores close to the standard baseline and is consistently 2\(\times\) faster. Can also shows a similar trend that it slightly underperforms the balanced baseline (\(<0.5\) BLEU scores) but is \(>27\%\) faster.

5 Analysis↩︎

5.1 Knowledge Distillation↩︎

Although San, Aan and Can offer considerable speed gain over the balanced baseline, they all suffer from the performance degradation as shown in Table [tab:wmt14] and Table [tab:wmt17]. The popular solution to this is knowledge distillation (KD). Here we choose sequence-level knowledge distillation [20] for better performance in machine translation tasks. The balanced baseline is used to generate the pseudo data for KD.

Table [tab:kd] shows that KD closes the performance gap between the fast attention models (San, Aan and Can) and the balanced baseline. This fact suggests that all three systems have enough capacity for a good performance, but training from scratch is not able to reach a good convergence state. It suggests that these systems might require a more careful hyper-parameters tuning or a better optimization method.

5.2 Ablation Study↩︎

To investigate in which part Can contributes the most to the acceleration as well as the performance loss, we only compress the self-attention and cross-attention or compress the cross-attention and FFN for study. Table 1 shows the results of this ablation study. We can see that compressing the two attentions provides a 20.09% speed-up, while only 6.21% for compressing attention and FFN. This is because FFN is already highly parallelized and accelerating itself does not bring much gain. Note that the second row of Table 1 is exactly [21]’s work (without AAN). We see that Can is faster and performs better. On the other hand, compressing attentions brings the most performance loss, which shows that the identical input assumption is strong. Fig. 2 shows that inputs of the adjacent layers are not very similar in lower layers. Therefore using Can in low layers might bring a great loss. We also find that compressing attention and FFN has an even better result. This might be that we remove the redundant parameters in the model.

5.3 Sensitivity Analysis↩︎

We study how the speed could be affected by other factors in Fig. 5, e.g., the beam size and the translation length. The left of Fig. 5 shows that Can is consistently faster than the balanced baseline with different beam size. As the acceleration provided by Can is constantly proportional to the speed of the baseline, it becomes less obvious when the baseline is slow, i.e., translating with a large beam. An opposite trend happens in the middle of Fig. 5 for the translation length. This is because overheads such as data preparation dominate the translation time of short sentences. This way results in a slow speed even when the translation time is short. As both the baseline and Can get faster when generating longer translations, one might suspect that the superior acceleration of Can over other approaches comes from the fact that Can generates longer translations. Further analysis is conducted and shown in the right of Fig. 5. We see that Can, San and Aan generate translations with similar lengths in two WMT14 translation tasks. This observation justifies that the superior acceleration brought by Can did come from its design rather than translation lengths.

5.4 Error Analysis↩︎

As shown in Table [tab:wmt14] and Table [tab:wmt17], the acceleration of Can comes at the cost of performance. Here we conduct experiments to better understand in which aspect Can scarifies for speed-up. We first evaluate the sentence-level BLEU score for each translation, then cluster these translations according to their averaged word frequencies or lengths.

Fig. 6 shows the results. The left of Fig. 6 indicates that Can did well on sentences with low frequencies, but not on those with high frequencies. The right of Fig. 6 shows that Can does not translate short sentences well but is quite good at translating long sentences. These facts are counterintuitive as one might expect a poor model could do well on easy samples (high frequency and short sentences) but not on hard ones (low frequency and long sentences). This might due to the identical input assumptions we used to derive Can are critical to easy samples. We left this for the future exploration.

5.5 Parallelism Study↩︎

A simple approach to obtain a higher parallelism without modifying the architecture is to increase the batch size at inference. Fig. 7 compares the inference time of the balanced baseline and Can by varying the batch size. We can see that both systems run faster with a larger batch size and Can is consistently faster than the balanced baseline. But the acceleration of Can over the baseline \(\Delta_{\mathrm{Speed}}\) diminishes when the batch size gets larger. In this case we observe that Can reaches the highest parallelism (a nearly 100% GPU utility) in a smaller batch size (\(\geq 32\)) than the baseline (\(>64\)). This means that enlarging the batch size no longer provides acceleration for Can, while the baseline can still be further speeded up. We expect Can could be faster if more tensor cores are available in the future.

5.6 Training Study↩︎

We plot the training and validation loss curve of the standard baseline, the balanced baseline and Can in Fig. 8 for studying their convergence. We can see that all systems converge stably. The balanced baseline has a higher loss than the standard baseline in both the training and validation sets, but their BLEU scores are close as shown in Table [tab:wmt14]. This is due to the shallow decoder in the balanced baseline. Since the loss is determined by the decoder, a shallow decoder with less capacity would have a higher loss. [2] indicates that the encoder depth has a greater impact than the decoder on BLEU scores, therefore the deep encoder makes up the performance loss of the shallow decoder. We also see that Can has a higher loss than the balanced baseline because we compress the decoder. Since we do not enhance the encoder, the BLEU score drops accordingly.

6 Related Work↩︎

6.1 Model Acceleration↩︎

Large Transformer has demonstrated its effectiveness on various natural language processing tasks, including machine translation [1], language modelling [22] and etc. The by-product brought by this huge network is the slow inference speed. Previous work focuses on improving model efficiency from different perspectives. For example, knowledge distillation approaches treat the large network output as the ground truth to train a small network [20]. Low-bit quantization approaches represent and run the model with 8-bit integer [23]. Our work follows another line of researches, which purses a more efficient architecture.

[24] show that the attention of Transformer benefits the encoder the most and the decoder could be safely replaced by a recurrent network. This way reduces the complexity of the decoder to linear time but incurs a high cost in training. [7] show that the self-attention is not necessary and a simple averaging is enough. [8] indicate that most attention distributions are redundant and thus share these distributions among layers. [25] use locality-sensitive hashing to select a constant number of words and perform attention on them. [26] train a large Transformer and drop some layers at testing for fast inference. [9] use a non-autoregressive decoder to predict the whole sentence at one time instead of generating it word by word. This approach makes a linear time translation process to constant time via the parallel computation. Perhaps the most related works are [21], [27]. They merge the self-attention and cross-attention and share their parameters. We, on the other hand, use different sets of parameters for each attention and mathematically prove that this way is equivalent to the standard Transformer under some mild conditions. We further show that the attention and FFN could also be merged together due to their linearity.

6.2 Deep Transformer↩︎

Recent studies have shown that deepening the Transformer encoder is more beneficial than widening the encoder or deepening the decoder [28]. [2] show that placing the layer normalization before (Pre-Norm) rather than behind (Post-Norm) the sub-layer allows us to train deep Transformer. [29] prove that the success of the Pre-Norm network relies on its well-behaved gradient. [21] suggest that a proper initialization is enough to train a deep Post-Norm network. [6] similarly exploit this observation but to build a faster instead of a better model. They show that using knowledge distillation, a deep encoder and shallow decoder model could run much faster without losing any performance. Based on their work, we use this model as our baseline system and evaluate it on extensive machine translation tasks without knowledge distillation.

7 Conclusion↩︎

In this work, we propose Can, whose decoder layer consists of only one attention. Can offers consistent acceleration by providing a high degree of parallelism. Experiments on 14 WMT machine translation tasks show that Can is 2.82\(\times\) faster than the baseline. We also use a stronger baseline for comparison. It employs a deep encoder and a shallow decoder, and is 2\(\times\) faster than the standard Transformer baseline without loss in performance.

Acknowledgments↩︎

This work was supported in part by the National Science Foundation of China (Nos. 61876035 and 61732005), and the National Key R&D Program of China (No. 2019QY1801). The authors would like to thank anonymous reviewers for their comments.

References↩︎