Speeding up Transformer Decoding via an Attention Refinement Network (2022)

Abstract: Despite the revolutionary advances made by Transformer in Neural Machine Translation (NMT), inference efficiency remains an obstacle due to the heavy use of attention operations in auto-regressive decoding. We thereby propose a lightweight attention structure called Attention Refinement Network (ARN) for speeding up Transformer. Specifically, we design a weighted residual network, which reconstructs the attention by reusing the features across layers. To further improve the Transformer efficiency, we merge the self-attention and cross-attention components for parallel computing. Extensive experiments on ten WMT machine translation tasks show that the proposed model yields an average of 1.35x faster (with almost no decrease in BLEU) over the state-of-the-art inference implementation. Results on widely used WMT14 En-De machine translation tasks demonstrate that our model achieves a higher speed-up, giving highly competitive performance compared to AAN and SAN models with fewer parameter numbers.

## introduction transformer @xcite has become the dominant approach in the nmt literature, which achieves superior translation performance and efficiency due to its well-designed attention mechanism. the highly parallelizable architecture enables transformer to capture the dependency among positions over the entire sequence parallelly for a faster training step. however, the inference efficiency remains a bottleneck for transformer. in inference, transformer follows an auto-regressive generation paradigm and generates the target words one by one on the decoder side. the heavy use of dot-product attention operations even further slows transformer efficiency. in addition, there are a figure 1 : the information transformation between the input and output after through a sub-layer f, where figure 2 : the cosine similarity of the input and output for every sub-layer over validation set on wmt14 @xmath0. "enc", "dec", "sa", "ca", "ffn" represent encoder, decoder, self-attention, cross-attention and feed-forward network, respectively. darker cells denote more similar. many efforts have been dedicated to accelerating the decoding process of transformer. @xcite adopts an average strategy to avoid computing the correlations over the entire input word. however, this method requires a complicated network and only focuses on the decoder-side self-attention. san @xcite reuses the attention results among layers, but it requires a model to learn which layers should be allowed to share. besides, another representative approach @xcite @xcite follows another line to abandon the auto-regressive generation property and produces target sequences in parallel, thus it fails to model the word dependencies. we investigate an alternative lightweight attention structure for transformer acceleration. as shown in figure 3 , the main idea is to reconstruct the attention via a weighted combination of highlevel and low-level features, rather than the standard dot-product function. the two parts represent the attention results of the current layer and its previous layer, respectively. this weight, a learned matrix under the guidance of the two input features, which determines how many low-level features can be selected to fuse. consequently, we refer to our model as attention refinement network, or arn in short. our combination process is computationally inexpensive, which allows us to employ less computation to obtain the attention results of adjacent layers and combine them to approximate the original attention results. arn structure can be viewed as a weighted residual network, which acquires the attention results by reusing the features across layers. this way introduces low-level features into the attention sub-layer to further enhance model confidence, which is an improvement over san @xcite . in addition, this method can be applied to both self-attention and cross-attention on the decoder side and we merge the above two components to further improve decoding efficiency. moreover, arn structure is simple and it is easy to implement. as another "bonus", arn requires fewer parameters, so it is faster to train and maintains a smaller memory footprint. extensive experiments on ten wmt translation tasks show that arn achieves an average of 1.35× faster with performance on par with a strong baseline. compared to aan and san baselines, our model gives a higher speed-up as well as highly competitive performance with fewer parameter numbers on widely used wmt14 @xmath1. ## standard transformer attention standard transformer follows the popular encoderdecoder paradigm, which consists of a 6-layer encoder and a 6-layer decoder. the overall architecture only contains stacked attention and feedforward networks (ffn) in transformer. there are three types of attention mechanisms: the encoder-side self-attention, the decoder-side self-attention and the cross-attention. on the encoder side, each layer follows the order of operations that could be defined as: self-@xmath2. similarly, the decoder side follows the way: self-@xmath3. the attention model in transformer is scaled dot-product attention. see figure 3 (a) for an illustration of the standard attention. the input of attention is a tuple of (q l , k l , v l ), where q @xmath4, v @xmath5, keys and values of the l-th layer. for encoder or decoder self-attention, @xmath6. for cross-attention, m and n are the target sequence length and source sequence length, respectively. d is the dimension of the hidden representation. we first compute the attention distribution via a scaled dot-product and softmax operations. ## attention refinement network (arn) the proposed method for speeding up transformer concentrates on the decoder-side attention because the decoder is the bottleneck of transformer inference. overall architecture. the proposed arn module is shown in figure 3 (b) . we assume that the decoder contains @xmath7. the decoder is simply divided into m parts equally, and each part contains n layers. for each arn module, the bottom layer is the dot-product attention same as used in standard transformer, and the next n -1 layers are the lightweight attention composed of the two inputs. one is an approximate attention result of the current layer by reusing the attention weights within adjacent layers, and the other is the up-sampled attention results from its previous layer. since the upper layer contains more semantic information, we refer to the two inputs high-level features and low-level features, respectively. α is a learned weight matrix, which is used to represent the fusion degree of low-level features. in contrast to figure 3 (a), arn reconstructs the attention via a feature fusion of the two input sources to replace the original dot-product attention. arn can be regarded as a weighted residual network, which sums to obtain the attention results by reusing the features across layers. the arn module is applied to both self-attention and cross-attention for speeding up transformer. ## analysis analysis on convergency. as shown in figure 5 , we plot the loss curves of standard transformer and arn model on the training and validation sets, respectively. the two systems converge stably and both of their loss curves maintain a high degree of similarity. they also present similar performance in bleu as shown in ## related work standard transformer suffers from the high inference cost due to the auto-regressive generation schema and the heavy use of dot-product attention operations. a classic solution is to generate the entire target sequence at one time by using non-autoregressive inference method @xcite @xcite . this way offers high decoding efficiency but it is hard to train. another representative solution adopts a local attention strategy @xcite or simplifies attention structure @xcite , which can effectively reduce the computation of attention module. but these works are designed for acceleration of very long sequence tasks (e.g., image generation, automatic speech recognition, etc.). nmt inference acceleration methods have been investigated for years, including knowledge distillation @xcite @xcite , vocabulary selection (l @xcite @xcite , low-precision computation @xcite @xcite , kernel fusion @xcite , layerdrop @xcite and etc. compared with the above methods, our work follows another line of work to learn a lightweight attention nmt model and prove its effectiveness. @xcite show that the self-attention network is not necessary and a simple averaging is enough. compared with the above method, our network structure is simple and it is easy to implement. we improve both the self-attention and cross-attention components. @xcite observe that the most attention distributions are similar and thus share these distributions among layers. this method can be regarded as a structure pruning for transformer, which may lead to model performance degradation. our approach, which fuses low-level features to further enhance model confidence, is an improvement over this method. @xcite propose the compressed attention network that simplifies the transformer architecture to achieve a higher parallelism. this method can be regarded as an equivalent transformation of the transformer structure, it is orthogonal to our research. arn structure can be viewed as a weighted residual network, which acquires the attention results by reusing the features across layers. the main idea is similar to @xcite , but our arn is different in both motivation and network structure. to our knowledge, we are the first to design a residual network for nmt decoding acceleration, proposing the combination of the features across layers for reconstructing the attention. ## conclusion we have investigated arn, an alternative lightweight attention structure for faster inference of transformer. experiments on a range of wmt translation tasks show that arn offers a significant speed improvement over a strong transformer baseline without sacrificing translation performance. results on widely used wmt14 @xmath8, compared to previous work like aan and san.