RealFormer: Transformer Likes Residual Attention

Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, pages 929–943 August 1–6, 2021. ©2021 Association for Computational Linguistics 929 RealFormer: Transformer Likes Residual Attention Ruining He, Anirudh Ravula, Bhargav Kanagal, Joshua Ainslie Google Research {ruininghe,braineater,bhargav,jainslie}@google.com Abstract Transformer is the backbone of modern NLP models. In this paper, we propose Real- Former, a simple and generic technique to create Residual Attention Layer Transformer networks that significantly outperform the canonical Transformer and its variants (BERT, ETC, etc.) on a wide spectrum of tasks including Masked Language Modeling, GLUE, SQuAD, Neural Machine Translation, WikiHop, HotpotQA, Natural Questions, and OpenKP. We also observe empirically that RealFormer stabilizes training and leads to models with sparser attention. Source code and pre-trained checkpoints for RealFormer can be found at https://github.com/ google-research/google-research/ tree/master/realformer. 1 Introduction Transformer (Vaswani et al., 2017) architectures are the backbone of numerous state-of-the-art NLP models such as BERT (Devlin et al., 2019), GPT (Radford et al., 2019), and Meena (Adiwar- dana et al., 2020), and have seen wide success across both academia and industry. Typically, a Transformer network consists of a stack of residual layers. The original design follows a “Post-LN” structure which adds Layer Norm (LN) as a “post- processing” step for each sub-layer, as shown in Figure 1 (a). It has been adopted by various state-of- the-art models including BERT, XLNet (Yang et al., 2019), RoBERTa (Liu et al., 2019), ALBERT (Lan et al., 2019), Transformer-XL (Dai et al., 2019), and ETC (Ainslie et al., 2020). Another notable design is to reorganize the order of modules to create a “direct”/clean path to propagate embed- dings of tokens in the input sequence through the whole network, as shown in Figure 1 (b).1 This 1Note that a final LN module is usually added at the very top of the whole network. design adds LN as a “pre-processing” step for each sub-layer, and is often referred to as “Pre-LN” and used by some well-known extra large models such as GPT-2 (Radford et al., 2019) and Mega- tron (Shoeybi et al., 2019). In some respect, Post- LN and Pre-LN are analogous to ResNet v1 (He et al., 2016a) and ResNet v2 (He et al., 2016b) respectively in the Computer Vision literature. Al- though ResNet v2 is usually preferable to v1 for Computer Vision, it does not appear to be the case for Pre-LN Transformer in the NLP literature. It is likely that the particularities of self-attention modules and Transformer architectures potentially favor (at least slightly) different designs compared to traditional convolutional neural networks. In this paper, we propose a simple and generic technique to show that it is beneficial to create a “di- rect” path to propagate raw attention scores through Transformer-based networks. Our technique is called Residual Attention Layer Transformer, or RealFormer in short. We also use RealFormer to denote the resulting Transformer networks when- ever no confusion may arise. Without losing gener- ality, taking the standard Transformer encoder as an example, each RealFormer layer takes the raw attention scores of all attention heads …

Content truncated to 500 words