Sharing Attention Weights for Fast Transformer

By Tong Xiao, Yinqiao Li, Jingbo Zhu, Zhengtao Yu, Tongran Liu. 2019

62 citations recorded by the source corpus.

Recently, the Transformer machine translation system has shown strong results by stacking attention layers on both the source and target-language sides. But the inference of this model is slow due to the heavy use of dot-product attention in auto-regressive decoding. In this paper we speed up Transformer via a fast and lightweight attention model. More specifically, we share attention weights in adjacent layers and enable the efficient re-use of hidden states in a vertical manner. Moreover, the sharing policy can be jointly learned with the MT model. We test our approach on ten WMT and NIST OpenMT tasks. Experimental results show that it yields an average of 1.3X speed-up (with almost no decrease in BLEU) on top of a state-of-the-art implementation that has already adopted a cache for fast inference. Also, our approach obtains a 1.8X speed-up when it works with the AAN model. This is even 16 times faster than the baseline with no use of the attention cache.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures, Efficiency & Compression, Machine Translation

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

In Efficiency & Compression

Citations in the corpus

Rethinking Attention with Performers (2020, neutral): “, 2020) and shared attention weights (Xiao et al., 2019).”
Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers (2020, background): “include reversible residual layers allowing for one-time activation storage in training [29] and shared attention weights [53].”
An Efficient Transformer Decoder with Compressed Sub-layers (2021, unknown): “S AN share the attention re-sults among layers (Xiao et al. 2019).”
The Cascade Transformer: an Application for Efficient Answer Sentence Selection (2020, methodology): “Being aware of this problem, the research community has been developing techniques to prune unnecessary network parameters (Lan et al., 2019; Sanh et al., 2019) or optimize the transformer architecture (Zhang et al., 2018; Xiao et al., 2019).”
Learning Light-Weight Translation Models from Deep Transformer (2020, result): “Similar findings were observed in previous work (Zhang, Titov, and Sennrich 2019; Xiao et al. 2019), which could speed up the Transformer inference by simplifying the decoder architecture.”
Audio Albert: A Lite Bert for Self-Supervised Learning of Audio Representation (2020, methodology): “Inspired by Fast Transformer [20] Experiments, we use the Jensen-Shannon (JS) divergence to evaluate the difference between the attention distribution of each transformer encoder layer.”
Towards Fully 8-bit Integer Inference for the Transformer Model (2020, methodology): “Also, the low-precision approach is orthogonal to other existing compression and acceleration methods, e.g., efﬁcient network design [ Xiao et al. , 2019 ] .”