Neural networks have been successfully applied to a variety of tasks, including machine translation. The encoder-decoder architecture is the central idea of neural machine translation (NMT). The encoder first encodes a source-side sentence into hidden states and then the decoder generates the target-side sentence from the hidden states according to an autoregressive model
Recurrent neural networks (RNNs) are inherently good at processing sequential data. sutskever2014sequence,cho2014learning successfully applied RNNs to machine translation. bahdanau2014neural introduced attention mechanism into the encoder-decoder architecture and greatly improved NMT. GNMT Wu et al. (2016)
The sequential property of RNNs leads to its wide application in language processing. However, the property also hinders its parallelizability thus RNNs are slow to execute on modern hardware optimized for parallel execution. As a result, a number of more parallelizable sequence models were proposed such as ConvS2S Gehring et al. (2017) and the Transformer Vaswani et al. (2017). These models avoid the dependencies between different positions in each layer thus can be trained much faster than RNN based models. When inference, however, these models are still slow because of the autoregressive property.
A recent work Gu et al. (2017) proposed a non-autoregressive NMT model that generates all target-side words in parallel. While the parallelizability is greatly improved, the translation quality encounter much decrease. In this paper, we propose the semi-autoregressive Transformer (SAT) for faster sequence generation. Unlike gu2017non, the SAT is semi-autoregressive, which means it keeps the autoregressive property in global but relieves in local. As the result, the SAT can produce multiple successive words in parallel at each time step. Figure 1 gives an illustration of the different levels of autoregressive properties.
Experiments conducted on English-German and Chinese-English translation show that compared with non-autoregressive methods, the SAT achieves a better balance between translation quality and decoding speed. On WMT’14 English-German translation, the proposed SAT is 5.58 faster than the Transformer while maintaining 88% of translation quality. Besides, when produces two words at each time step, the SAT is almost lossless.
It is worth noting that although we apply the SAT to machine translation, it is not designed specifically for translation as gu2017non,lee2018deterministic. The SAT can also be applied to any other sequence generation task, such as summary generation and image caption generation.
2 Related Work
Almost all state-of-the-art NMT models are autoregressive Sutskever et al. (2014); Bahdanau et al. (2014); Wu et al. (2016); Gehring et al. (2017); Vaswani et al. (2017), meaning that the model generates words one by one and is not friendly to modern hardware optimized for parallel execution. A recent work Gu et al. (2017) attempts to accelerate generation by introducing a non-autoregressive model. Based on the Transformer Vaswani et al. (2017), they made lots of modifications. The most significant modification is that they avoid feeding the previously generated target words to the decoder, but instead feeding the source words, to predict the next target word. They also introduced a set of latent variables to model the fertilities
of source words to tackle the multimodality problem in translation. lee2018deterministic proposed another non-autoregressive sequence model based on iterative refinement. The model can be viewed as both a latent variable model and a conditional denoising autoencoder. They also proposed a learning algorithm that is hybrid of lower-bound maximization and reconstruction error minimization.
The most relevant to our proposed semi-autoregressive model is Kaiser et al. (2018). They first autoencode the target sequence into a shorter sequence of discrete latent variables, which at inference time is generated autoregressively, and finally decode the output sequence from this shorter latent sequence in parallel. What we have in common with their idea is that we have not entirely abandoned autoregressive, but rather shortened the autoregressive path.
A related study on realistic speech synthesis is the parallel WaveNet Oord et al. (2017). The paper introduced probability density distillation, a new method for training a parallel feed-forward network from a trained WaveNet Van Den Oord et al. (2016) with no significant difference in quality.
There are also some work share a somehow simillar idea with our work: character-level NMT Chung et al. (2016); Lee et al. (2016) and chunk-based NMT Zhou et al. (2017); Ishiwatari et al. (2017). Unlike the SAT, these models are not able to produce multiple tokens (characters or words) each time step. oda2017neural proposed a bit-level decoder, where a word is represented by a binary code and each bit of the code can be predicted in parallel.
3 The Transformer
Since our proposed model is built upon the Transformer Vaswani et al. (2017), we will briefly introduce the Transformer. The Transformer uses an encoder-decoder architecture. We describe the encoder and decoder below.
3.1 The Encoder
From the source tokens, learned embeddings of dimension are generated which are then modified by an additive positional encoding. The positional encoding is necessary since the network does not leverage the order of the sequence by recurrence or convolution. The authors use additive encoding which is defined as:
where is the position of a word in the sentence and is the dimension. The authors chose this function because they hypothesized it would allow the model to learn to attend by relative positions easily. The encoded word embeddings are then used as input to the encoder which consists of blocks each containing two layers: (1) a multi-head attention layer, and (2) a position-wise feed-forward layer.
Multi-head attention builds upon scaled dot-product attention, which operates on a query Q, key K and value V:
where is the dimension of the key. The authors scale the dot product by to avoid the inputs to softmax function growing too large in magnitude. Multi-head attention computes different queries, keys and values with linear projections, computes scaled dot-product attention for each query, key and value, concatenates the results, and projects the concatenation with another linear projection:
in which and . The attention mechanism in the encoder performs attention over itself (), so it is also called self-attention.
The second component in each encoder block is a position-wise feed-forward layer defined as:
where , , , .
3.2 The Decoder
The decoder is similar with the encoder and is also composed by blocks. In addition to the two layers in each encoder block, the decoder inserts a third layer, which performs multi-head attention over the output of the encoder.
It is worth noting that, different from the encoder, the self-attention layer in the decoder must be masked with a causal mask, which is a lower triangular matrix, to ensure that the prediction for position can depend only on the known outputs at positions less than during training.
4 The Semi-Autoregressive Transformer
We propose a novel NMT model—the Semi-Autoregressive Transformer (SAT)—that can produce multiple successive words in parallel. As shown in Figure 2, the architecture of the SAT is almost the same as the Transformer, except some modifications in the decoder.
4.1 Group-Level Chain Rule
Standard NMT models usually factorize the joint probability of a word sequence
according to the word-level chain rule
resulting in decoding each word depending on all previous decoding results, thus hindering the parallelizability. In the SAT, we extend the standard word-level chain rule to the group-level chain rule.
We first divide the word sequence into consecutive groups
where denotes floor operation, is the group size, and also the indicator of parallelizability. The larger the , the higher the parallelizability. Except for the last group, all groups must contain words. Then comes the group-level chain rule
This group-level chain rule avoids the dependencies between consecutive words if they are in the same group. With group-level chain rule, the model no longer produce words one by one as the Transformer, but rather group by group. In next subsections, we will show how to implement the model in detail.
4.2 Long-Distance Prediction
In autoregressive models, to predict , the model should be fed with the previous word . We refer it as short-distance prediction. In the SAT, however, we feed to predict , to which we refer as long-distance prediction. At the beginning of decoding, we feed the model with special symbols to predict in parallel. Then are fed to the model to predict in parallel. This process will continue until a terminator is generated. Figure 3 gives illustrations for both short and long-distance prediction.
4.3 Relaxed Causal Mask
In the Transformer decoder, the causal mask is a lower triangular matrix, which strictly prevents earlier decoding steps from peeping information from later steps. We denote it as strict causal mask. However, in the SAT decoder, strict causal mask is not a good choice. As described in the previous subsection, in long-distance prediction, the model predicts by feeding with . With strict causal mask, the model can only access to when predict , which is not reasonable since are already produced. It is better to allow the model to access to rather than only when predict .
Therefore, we use a coarse-grained lower triangular matrix as the causal mask that allows peeping later information in the same group. We refer to it as relaxed causal mask. Given the target length and the group size , relaxed causal mask and its elements are defined below:
For a more intuitive understanding, Figure 4 gives a comparison between strict and relaxed causal mask.
4.4 The SAT
Using group-level chain rule instead of word-level chain rule, long-distance prediction instead of short-distance prediction, and relaxed causal mask instead of strict causal mask, we successfully extended the Transformer to the SAT. The Transformer can be viewed as a special case of the SAT, when the group size = 1. The non-autoregressive Transformer (NAT) described in gu2017non can also be viewed as a special case of the SAT, when the group size is not less than maximum target length.
Table 1 gives the theoretical complexity and acceleration of the model. We list two search strategies separately: beam search and greedy search. Beam search is the most prevailing search strategy. However, it requires the decoder states to be updated once every word is generated, thus hinders the decoding parallelizability. When decode with greedy search, there is no such concern, therefore the parallelizability of the SAT can be maximized.
|SAT (beam search)|
|SAT (greedy search)|
We evaluate the proposed SAT on English-German and Chinese-English translation tasks.
5.1 Experimental Settings
Datasets For English-German translation, we choose the corpora provided by WMT 2014 Bojar et al. (2014). We use the newstest2013 dataset for development, and the newstest2014 dataset for test. For Chinese-English translation, the corpora we use is extracted from LDC111The corpora include LDC2002E18, LDC2003E14, LDC2004T08 and LDC2005T0.. We chose the NIST02 dataset for development, and the NIST03, NIST04 and NIST05 datasets for test. For English and German, we tokenized and segmented them into subword symbols using byte-pair encoding (BPE) Sennrich et al. (2015) to restrict the vocabulary size. As for Chinese, we segmented sentences into characters. For English-German translation, we use a shared source and target vocabulary. Table 2 summaries the two corpora.
|Sentence Number||Vocab Size|
|NAT Gu et al. (2017)||-||17.69||25%||39ms||15.6|
|NAT (rescroing 10)||-||18.66||20%||79ms||7.68|
|NAT (rescroing 100)||-||19.17||18%||257ms||2.36|
|LT Kaiser et al. (2018)||-||19.80||27%||105ms||-|
|LT (rescoring 10)||-||21.00||23%||-||-|
|LT (rescoring 100)||-||22.50||18%||-||-|
|IRNAT Lee et al. (2018)||-||18.91||22%||-||1.98|
Results on English-German translation. Latency is calculated on a single NVIDIA TITAN Xp without batching. For comparison, we also list results reported by gu2017non,kaiser2018fast,lee2018deterministic. Note that gu2017non,lee2018deterministic used PyTorch as their platform, but we and kaiser2018fast used TensorFlow. Even on the same platform, implementation and hardware may not exactly be the same. Therefore, it is not fair to directly compare BLEU and latency. A fairer way is to compare performance degradation and speedup, which are calculated based on their own baseline.
Baseline We use the base Transformer model described in vaswani2017attention as the baseline, where . In addition, for comparison, we also prepared a lighter Transformer model, in which two encoder/decoder blocks are used ( = 2), and other hyper-parameters remain the same.
Unless otherwise specified, all hyperparameters are inherited from the base Transformer model. We try three different settings of the group size: = 2, = 4, and = 6. For English-German translation, we share the same weight matrix between the source and target embedding layers and the pre-softmax linear layer. For Chinese-English translation, we only share weights of the target embedding layer and the pre-softmax linear layer.
Search Strategies We use two search strategies: beam search and greedy search. As mentioned in Section 4.4, these two strategies lead to different parallelizability. When beam size is set to 1, greedy search is used, otherwise, beam search is used.
Knowledge Distillation Knowledge distillation Hinton et al. (2015); Kim and Rush (2016) describes a class of methods for training a smaller student network to perform better by learning from a larger teacher
network. For NMT, kim2016sequence proposed a sequence-level knowledge distillation method. In this work, we apply this method to train the SAT using a pre-trained autoregressive Transformer network. This method consists of three steps: (1) train an autoregressive Transformer network (theteacher), (2) run beam search over the training set with this model and (3) train the SAT (the student) on this new created corpus.
Initialization Since the SAT and the Transformer have only slight differences in their architecture (see Figure 2), in order to accelerate convergence, we use a pre-trained Transformer model to initialize some parameters in the SAT. These parameters include all parameters in the encoder, source and target word embeddings, and pre-softmax weights. Other parameters are initialized randomly. In addition to accelerating convergence, we find this method also slightly improves the translation quality.
Training Same as vaswani2017attention, we train the SAT by minimize cross-entropy with label smoothing. The optimizer we use is Adam Kingma and Ba (2015) with , and . We change the learning rate during training using the learning rate funtion described in vaswani2017attention. All models are trained for 10K steps on 8 NVIDIA TITAN Xp with each minibatch consisting of about 30k tokens. For evaluation, we average last five checkpoints saved with an interval of 1000 training steps.
5.2 Results on English-German
Table 3 summaries results of English-German translation. According to the results, the translation quality of the SAT gradually decreases as increases, which is consistent with intuition. When = 2, the SAT decodes 1.51 faster than the Transformer and is almost lossless in translation quality (only drops 0.21 BLEU score). With = 6, the SAT can achieve 2.98 speedup while the performance degeneration is only 8%.
When using greedy search, the acceleration becomes much more significant. When = 6, the decoding speed of the SAT can reach about of the Transformer while maintaining 88% of translation quality. Comparing with gu2017non,kaiser2018fast,lee2018deterministic, the SAT achieves a better balance between translation quality and decoding speed. Compared to the lighter Transformer ( = 2), with = 4, the SAT achieves a higher speedup with significantly better translation quality.
In a real production environment, it is often not to decode sentences one by one, but batch by batch. To investigate whether the SAT can accelerate decoding when decoding in batches, we test the decoding latency under different batch size settings. As shown in Table 4, the SAT significantly accelerates decoding even with a large batch size.
It is also good to know if the SAT can still accelerate decoding on CPU device that does not support parallel execution as well as GPU. Results in Table 5 show that even on CPU device, the SAT can still accelerate decoding significantly.
5.3 Results on Chinese-English
Table 6 summaries results on Chinese-English translation. With = 2, the SAT decodes 1.69 while maintaining 97% of the translation quality. In an extreme setting where = 6 and beam size = 1, the SAT can achieve 6.41 speedup while maintaining 83% of the translation quality.
Effects of Knowledge Distillation As shown in Figure 5, sequence-level knowledge distillation is very effective for training the SAT. For larger , the effect is more significant. This phenomenon is echoing with observations by gu2017non,oord2017parallel,lee2018deterministic. In addition, we tried word-level knowledge distillation Kim and Rush (2016) but only a slight improvement was observed.
Position-Wise Cross-Entropy In Figure 6, we plot position-wise cross-entropy for various models. To compare with the baseline model, the results in the figure are from models trained on the original corpora, i.e., without knowledge distillation. As shown in the figure, position-wise cross-entropy has an apparent periodicity with a period of . For positions in the same group, the position-wise cross-entropy increase monotonously, which indicates that the long-distance dependencies are always more difficult to model than short ones. It suggests the key to further improve the SAT is to improve the ability of modeling long-distance dependencies.
Case Study Table 7 lists three sample Chinese-English translations from the development set. As shown in the table, even when produces = 6 words at each time step, the model can still generate fluent sentences. As reported by gu2017non, instances of repeated words or phrases are most prevalent in their non-autoregressive model. In the SAT, this is also the case. This suggests that we may be able to improve the translation quality of the SAT by reducing the similarity of the output distribution of adjacent positions.
In this work, we have introduced a novel model for faster sequence generation based on the Transformer Vaswani et al. (2017), which we refer to as the semi-autoregressive Transformer (SAT). Combining the original Transformer with group-level chain rule, long-distance prediction and relaxed causal mask, the SAT can produce multiple consecutive words at each time step, thus speedup decoding significantly. We conducted experiments on English-German and Chinese-English translation. Compared with previously proposed non-autoregressive models Gu et al. (2017); Lee et al. (2018); Kaiser et al. (2018), the SAT achieves a better balance between translation quality and decoding speed. On WMT’14 English-German translation, the SAT achieves 5.58 speedup while maintaining 88% translation quality, significantly better than previous methods. When produces two words at each time step, the SAT is almost lossless (only 1% degeneration in BLEU score).
In the future, we plan to investigate better methods for training the SAT to further shrink the performance gap between the SAT and the Transformer. Specifically, we believe that the following two directions are worth study. First, use object function beyond maximum likelihood to improve the modeling of long-distance dependencies. Second, explore new method for knowledge distillation. We also plan to extend the SAT to allow the use of different group sizes at different positions, instead of using a fixed value.
We would like to thank the anonymous reviewers for their valuable comments. We also thank Wenfu Wang, Hao Wang for helpful discussion and Linhao Dong, Jinghao Niu for their help in paper writting.
Abadi et al. (2016)
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey
Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al.
Tensorflow: A system for large-scale machine learning.In OSDI, volume 16, pages 265–283.
- Ba et al. (2016) Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. 2016. Layer normalization. arXiv preprint arXiv:1607.06450.
- Bahdanau et al. (2014) Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. 2014. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473.
- Bojar et al. (2014) Ondrej Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Koehn, Johannes Leveling, Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, et al. 2014. Findings of the 2014 workshop on statistical machine translation. In Proceedings of the ninth workshop on statistical machine translation, pages 12–58.
- Cho et al. (2014) Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. 2014. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint arXiv:1406.1078.
- Chung et al. (2016) Junyoung Chung, Kyunghyun Cho, and Yoshua Bengio. 2016. A character-level decoder without explicit segmentation for neural machine translation. arXiv preprint arXiv:1603.06147.
- Gehring et al. (2017) Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N Dauphin. 2017. Convolutional sequence to sequence learning.
- Gu et al. (2017) Jiatao Gu, James Bradbury, Caiming Xiong, Victor OK Li, and Richard Socher. 2017. Non-autoregressive neural machine translation. arXiv preprint arXiv:1711.02281.
- He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016. Deep residual learning for image recognition. In
- Hinton et al. (2015) Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. 2015. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531.
- Ishiwatari et al. (2017) Shonosuke Ishiwatari, Jingtao Yao, Shujie Liu, Mu Li, Ming Zhou, Naoki Yoshinaga, Masaru Kitsuregawa, and Weijia Jia. 2017. Chunk-based decoder for neural machine translation. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), volume 1, pages 1901–1912.
- Kaiser et al. (2018) Łukasz Kaiser, Aurko Roy, Ashish Vaswani, Niki Pamar, Samy Bengio, Jakob Uszkoreit, and Noam Shazeer. 2018. Fast decoding in sequence models using discrete latent variables. arXiv preprint arXiv:1803.03382.
- Kim and Rush (2016) Yoon Kim and Alexander M Rush. 2016. Sequence-level knowledge distillation. arXiv preprint arXiv:1606.07947.
- Kingma and Ba (2015) Diederik P Kingma and Jimmy Lei Ba. 2015. Adam: A method for stochastic optimization. international conference on learning representations.
- Lee et al. (2016) Jason Lee, Kyunghyun Cho, and Thomas Hofmann. 2016. Fully character-level neural machine translation without explicit segmentation. arXiv preprint arXiv:1610.03017.
- Lee et al. (2018) Jason Lee, Elman Mansimov, and Kyunghyun Cho. 2018. Deterministic non-autoregressive neural sequence modeling by iterative refinement. arXiv preprint arXiv:1802.06901.
- Oda et al. (2017) Yusuke Oda, Philip Arthur, Graham Neubig, Koichiro Yoshino, and Satoshi Nakamura. 2017. Neural machine translation via binary code prediction. arXiv preprint arXiv:1704.06918.
- Oord et al. (2017) Aaron van den Oord, Yazhe Li, Igor Babuschkin, Karen Simonyan, Oriol Vinyals, Koray Kavukcuoglu, George van den Driessche, Edward Lockhart, Luis C Cobo, Florian Stimberg, et al. 2017. Parallel wavenet: Fast high-fidelity speech synthesis. arXiv preprint arXiv:1711.10433.
- Papineni et al. (2002) Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. 2002. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting on association for computational linguistics, pages 311–318. Association for Computational Linguistics.
- Sennrich et al. (2015) Rico Sennrich, Barry Haddow, and Alexandra Birch. 2015. Neural machine translation of rare words with subword units. arXiv preprint arXiv:1508.07909.
- Srivastava et al. (2014) Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. 2014. Dropout: A simple way to prevent neural networks from overfitting. The Journal of Machine Learning Research, 15(1):1929–1958.
- Sutskever et al. (2014) Ilya Sutskever, Oriol Vinyals, and Quoc V Le. 2014. Sequence to sequence learning with neural networks. In Advances in neural information processing systems, pages 3104–3112.
- Van Den Oord et al. (2016) Aaron Van Den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu. 2016. Wavenet: A generative model for raw audio. arXiv preprint arXiv:1609.03499.
- Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. arXiv preprint arXiv:1706.03762.
- Wu et al. (2016) Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, and Klaus Macherey. 2016. Google’s neural machine translation system: Bridging the gap between human and machine translation.
- Zhou et al. (2017) Hao Zhou, Zhaopeng Tu, Shujian Huang, Xiaohua Liu, Hang Li, and Jiajun Chen. 2017. Chunk-based bi-scale decoder for neural machine translation. arXiv preprint arXiv:1705.01452.