Sparse Attention with Linear Units

Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 6507–6520 November 7–11, 2021. c⃝2021 Association for Computational Linguistics 6507 Sparse Attention with Linear Units Biao Zhang1 Ivan Titov1,2 Rico Sennrich3,1 1School of Informatics, University of Edinburgh 2ILLC, University of Amsterdam 3Department of Computational Linguistics, University of Zurich B.Zhang@ed.ac.uk, ititov@inf.ed.ac.uk, sennrich@cl.uzh.ch Abstract Recently, it has been argued that encoder- decoder models can be made more inter- pretable by replacing the softmax function in the attention with its sparse variants. In this work, we introduce a novel, simple method for achieving sparsity in attention: we replace the softmax activation with a ReLU, and show that sparsity naturally emerges from such a for- mulation. Training stability is achieved with layer normalization with either a specialized initialization or an additional gating function. Our model, which we call Rectified Linear Attention (ReLA), is easy to implement and more efficient than previously proposed sparse attention mechanisms. We apply ReLA to the Transformer and conduct experiments on five machine translation tasks. ReLA achieves translation performance comparable to several strong baselines, with training and decoding speed similar to that of the vanilla attention. Our analysis shows that ReLA delivers high sparsity rate and head diversity, and the in- duced cross attention achieves better accuracy with respect to source-target word alignment than recent sparsified softmax-based models. Intriguingly, ReLA heads also learn to attend to nothing (i.e. ‘switch off’) for some queries, which is not possible with sparsified softmax alternatives.1 1 Introduction Attention models (Bahdanau et al., 2015) have been hugely successful recently, with Trans- former (Vaswani et al., 2017) in particular, advanc- ing state of the art on various tasks, such as machine translation (Bojar et al., 2018), document summa- rization (Liu and Lapata, 2019) and speech process- ing (Chiu et al., 2018), and delivering a large im- pact on a broad range of NLP tasks via large-scale self-supervised pretraining (Devlin et al., 2019). 1Source code is available at https://github.com/ bzhangGo/zero. Linear Linear Linear query context MatMul Scale Softmax Linear MatMul Linear Linear Linear query context MatMul Scale ReLU LayerNorm MatMul Linear Q K V Q K V (a) Attention with Softmax Linear Linear Linear query context MatMul Scale Softmax Linear MatMul Linear Linear Linear query context MatMul Scale ReLU LayerNorm MatMul Linear Q K V Q K V (b) Rectified Linear Attention Figure 1: Overview of the vanilla dot-product attention with softmax and the proposed Rectified Linear Attention (ReLA). Major differences highlighted in red. At the core of attention is a mechanism that dy- namically highlights relevant context features for a given query input. In the vanilla softmax-based attention model (Vaswani et al., 2017, SMATT), this is achieved by imposing a categorical distribu- tion constraint on the query-context relevance (i.e. attention) scores, implemented with the softmax activation (see Figure 1(a)). SMATT produces dense distributions, assigning some small amounts of attention even to irrelevant features. This complicates the analysis of the infor- mation flow in the model, and has led researchers to study sparse alternatives, …

Content truncated to 500 words