Speeding up Transformer Decoding via an Attention Refinement Network

Proceedings of the 29th International Conference on Computational Linguistics, pages 5109–5118 October 12–17, 2022. 5109 Speeding Up Transformer Decoding via an Attention Refinement Network Kaixin Wu1, Yue Zhang2,3, Bojie Hu1, Tong Zhang1 1Tencent Minority-Mandarin Translation, Beijing, China 2School of Engineering, Westlake University 3Institute of Advanced Technology, Westlake Institute for Advanced Study {danielkxwu, bojiehu, zatozhang}@tencent.com zhangyue@westlake.edu.cn Abstract Despite the revolutionary advances made by Transformer in Neural Machine Translation (NMT), inference efficiency remains an obsta- cle due to the heavy use of attention operations in auto-regressive decoding. We thereby pro- pose a lightweight attention structure called Attention Refinement Network (ARN) for speed- ing up Transformer. Specifically, we design a weighted residual network, which reconstructs the attention by reusing the features across lay- ers. To further improve the Transformer effi- ciency, 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.35× faster (with al- most no decrease in BLEU) over the state-of- the-art inference implementation. Results on widely used WMT14 En→De machine transla- tion tasks demonstrate that our model achieves a higher speed-up, giving highly competitive performance compared to AAN and SAN mod- els with fewer parameter numbers1. 1 Introduction Transformer (Vaswani et al., 2017) 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 architec- ture 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 1https://github.com/ Kaixin-Wu-for-Open-Source/ARN F ⊕ Input Output Xl Xl Xl+1 Figure 1: The information transformation between the input and output after through a sub-layer F, where Xl+1 = Xl + F(Xl) Figure 2: The cosine similarity of the input and out- put for every sub-layer over validation set on WMT14 En→De translation task. “Enc”, “Dec”, “SA”, “CA”, “FFN” represent Encoder, Decoder, Self-attention, Cross-attention and Feed-forward network, respectively. Darker cells denote more similar. large number of attention sub-layers in the multi- layer stacked Transformer, which makes it prone to redundancy. As shown in Figure 2, we observe redundant computation in Transformer attention sub-layer, especially the attention in the decoder. This inspires us to explore a lightweight attention structure to speed up Transformer decoding. Many efforts have been dedicated to accelerat- ing the decoding process of Transformer. AAN (Zhang et al., 2018) adopts an average strategy to avoid computing the correlations over the en- tire input word. However, this method requires a complicated network and only focuses on the decoder-side self-attention. SAN (Xiao et al., 2019) reuses the attention results among layers, but it requires a model to learn which layers should be allowed to share. Besides, another representative 5110 approach …

Content truncated to 500 words