Multi-Head Highly Parallelized LSTM Decoder for Neural Machine Translation (2021)

Abstract: One of the reasons Transformer translation models are popular is that self-attention networks for context modelling can be easily parallelized at sequence level. However, the computational complexity of a self-attention network is O(n^2) , increasing quadratically with sequence length. By contrast, the complexity of LSTM-based approaches is only O(n). In practice, however, LSTMs are much slower to train than self-attention networks as they cannot be parallelized at sequence level: to model context, the current LSTM state relies on the full LSTM computation of the preceding state. This has to be computed n times for a sequence of length n. The linear transformations involved in the LSTM gate and state computations are the major cost factors in this. To enable sequence-level parallelization of LSTMs, we approximate full LSTM context modelling by computing hidden states and gates with the current input and a simple bag-of-words representation of the preceding tokens context. This allows us to compute each input step efficiently in parallel, avoiding the formerly costly sequential linear transformations. We then connect the outputs of each parallel step with computationally cheap element-wise computations. We call this the Highly Parallelized LSTM. To further constrain the number of LSTM parameters, we compute several small HPLSTMs in parallel like multi-head attention in the Transformer. The experiments show that our MHPLSTM decoder achieves significant BLEU improvements, while being even slightly faster than the self-attention network in training, and much faster than the standard LSTM.

## introduction the transformer translation model @xcite has achieved great success and is used extensively in the nlp community. it achieves outstanding performance compared to previous rnn/cnn based translation models @xcite while being much faster to train. the transformer can be trained efficiently due to the highly parallelized self-attention network. it enables sequence-level parallelization in context modelling, as all token representations can be computed in parallel, and linear transformations are only required to compute the sequence once. on the other hand, previous rnn-based methods process a sequence in a token-by-token manner, which means that they have to compute linear layers once for each token, i.e. n times if the number of tokens in the sequence is n. however, the complexity of a self-attention network which compares each token with all the other tokens is o(n 2 ), while for lstm @xcite it is only o(n). in practice, however, lstm is slower than the self-attention network in training. this is mainly due to the fact that the computation of its current step relies on the computation output of the previous step, which prevents efficient parallelization over the sequence. as for the performance of using recurrent models in machine translation, @xcite shows that an lstm-based decoder can further improve the performance over the transformer. in this paper, we investigate how we can efficiently parallelize all linear transformations of an lstm at the sequence level, i.e. compute its linear transformations only once with a given input sequence. given that linear transformations are implemented by matrix multiplication, compared to the other element-wise operations, we suggest that they take the largest part of the model's overall computation, and parallelizing the linear transformations at sequence level may significantly accelerate the training of lstm-based models. our contributions are as follows: σ σ σ ɑ i t o t-1 concat * c t-1 * + * f g t o g t i g t h t c t o t figure 1: lstm. layer normalization is omitted for simplicity. ## preliminaries: lstm we design our hplstm based on the layer normalization @xcite enhanced lstm (ln-lstm) presented by @xcite as illustrated in figure 1 , which achieves better performance than the transformer when used in decoding. for the computation of gates and the hidden state, the model concatenates the input i t of the current step t to the output of the previous step o t-1 : and bias parameters, σ indicates the sigmoid activation function, α is the activation function for the hidden state computation, ln is the layer normalization. layer normalization (ba et al., 2016) is computed as follows: ## related work sequence-to-sequence neural machine translation models started with recurrent models @xcite @xcite . but recurrent models cannot be parallelized at the sequence level. convolutional models @xcite and the transformer @xcite have been proposed. due to the o(n 2 ) self-attention network complexity, which slows down decoding, @xcite presented the average attention network to accelerate decoding. even though lstms cannot be parallelized at the sequence level, its complexity is o(n), and @xcite shows that using the layer normalization enhanced lstm-based decoder can bring improvements in translation quality and accelerate decoding. lstm @xcite and gru @xcite are the most popular recur-rent models. to accelerate rnn models, @xcite propose a heavily simplified atr network to have the smallest number of weight matrices among units of all existing gated rnns. @xcite investigate exponentially decaying bag-of-words input features for feedforward nmt models. in addition to sequencelevel parallelization, asynchronous optimization @xcite and data parallelization with a larger batch size @xcite @xcite can also accelerate training. ## conclusion in this paper, we observe that the sequence-level parallelization issue of lstm is due to the fact that its computation of gates and hidden states of the current step relies on the computation result of the preceding step, and linear transformations have to be propagated the same number of times as the sequence length. to improve the sequencelevel parallelization of the lstm, we propose to remove the dependency of the current step lstm computation on the result of the previous step by computing hidden states and gates with the current input embedding and a bag-of-words representation of preceding tokens, and present the highly parallelized lstm. to constrain the number of lstm parameters, we compute several small hplstms in parallel like multi-head self-attention. in our experiments, we empirically show that the mhplstm model achieves better performance than self-attention networks, while being even slightly faster in training, and much faster in decoding, than the self-attention transformer decoder.