Luna: Linear Unified Nested Attention

1 Introduction Transformers ( Vaswani et al. , 2017 ) are surprisingly versatile models that preform well on a wide range of language and vision tasks , including machine translation ( Vaswani et al. , 2017 ; Ott et al. , 2018 ) , language understanding ( Devlin et al. , 2019 ) , image recognition ( Dosovitskiy et al. , 2020 ) and bioinformatics ( Madani et al. , 2020 ) . Attention ( Bahdanau et al. , 2015 ) provides the key mechanism that captures contextual information from the entire sequence by modeling pairwise interactions between the inputs at every timestep . However , a common weakness of Transformers is their quadratic time and memory complexity within the attention mechanism w.r.t the length of the input sequence , which prohibitively restricts their potential application to tasks requiring longer input sequences . A number of techniques have been recently introduced to improve the time and memory efficiency of Transformer models ( ‘ xformers ’ ) ( Tay et al. , 2020b , 2021 ) . One popular technique is using sparsity to restrict the attention field range , such as local attention ( Parmar et al. , 2018 ) , blockwise attention ( Qiu et al. , 2019 ) , strided attention patterns ( Child et al. , 2019 ; Beltagy et al. , 2020 ) , ∗Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . compressed attention ( Liu et al. , 2018 ) , and attention with learnable patterns ( Kitaev et al. , 2020 ; Tay et al. , 2020a ; Roy et al. , 2021 ) . Another emerging approach is to improve efficiency by leveraging low-rank approximations of the attention matrix . Linformer ( Wang et al. , 2020 ) , for example , projects the length dimension of key and value matrices to a fixed-dimensional representation by assuming low-rank structure in the full-rank attention matrix . Recently , some kernel-based methods , such as Linear Transformer ( Katharopoulos et al. , 2020 ) , Performer ( Choromanski et al. , 2020 ) and Random Feature Attention ( Peng et al. , 2021 ) , attempt to efficiently approximate regular ( softmax ) full-rank attention through kernelization . Although these models demonstrate better asymptotic complexity for long sequences , their efficiency gains are less prominent for moderate length sequences and their performance remains behind Transformers with regular attention . In this work , we propose a linear unified nested attention mechanism ( Luna ) , which uses two nested attention functions to approximate the regular softmax attention in Transformer ( §2 ) . Specifically , with the first attention function , Luna packs the input sequence into a sequence of fixed length . Then , the packed sequence is unpacked using the second attention function ( §3.1 ) . As compared to a more traditional attention mechanism , Luna introduces an additional sequence with a fixed length as input and an additional corresponding output . Importantly , the extra input allows Luna to perform attention operation linearly as efficiently as Linformer ( Wang et al. , 2020 ) , while also storing adequate contextual information . Unlike Linformer , Luna is capable of modeling variable-length sequences and autoregressive ( causal ) attention ( §3.4 ) . We perform extensive experiments on three sequence modeling tasks , including long-context sequence modeling , neural machine translation , and masked language modeling for large-scale pretraining and downstream task finetuning . Compared to a variety of strong baseline models , Luna achieves competitive or even better performance , while acquiring prominent gains of efficiency in both speed and memory ( see Figure 1 ) . More importantly , Luna manages to obtain superior performance with small projection lengths such as 16 ( §4 ) . 2 Background . 2.1 Attention . The traditional attention mechanism is a function : Y = Attn ( X , C ) = ω ( XWQ ( CWK ) T √ d ) CWV ( 1 ) where the attention function Attn : Rn×d×Rm×d → Rn×d takes as inputs two sequences : the query sequence X ∈ Rn×d with length n and the context sequence C ∈ Rm×d with length m , and output one sequence Y ∈ Rn×d with the same length n as the query X . d is the embedding dimension , and WQ , WK , WV ∈ Rd×d are three learnable parameters that project the input sequences into the space of query , key and value matrices : Q = XWQ , K = CWK , V = CWV . ω is an activation function , e.g . the softmax function in regular attention . Note that the formulation in ( 1 ) is applicable to both cross-attention where C andX are the representations from Transformer encoder and decoder , respectively , and self-attention where X and C are the same sequence ( X = C ) . In practice , the multi-head variant of attention ( Vaswani et al. , 2017 ) , which performs the attention function h times in parallel , is commonly used . Throughout this paper , we omit h for simplicity . In particular , the matrix A = ω ( QK T √ dk ) ∈ Rn×m in ( 1 ) is called the attention matrix which specifies the alignment scores between every pair of tokens in sequences of queries X and contexts C. Calculating A takes O ( nm ) time and space , which is quadratic with respect to the sequence length and becomes a significant bottleneck when processing long sequences . 2.2 Transformer Layers . The other two key components of Transformer , besides attention , are position-wise feed-forward networks ( FFN ) and layer normalization ( Ba et al. , 2016 ) . Technically , the position-wise feedforward layer operates on each position independently and layer normalization plays a crucial role in controlling the gradient scales ( Xiong et al. , 2020 ) . Each Transformer layer can be expressed as : XA = LayerNorm ( Attn ( X , C ) +X ) X ′ = LayerNorm ( FFN ( XA ) +XA ) ( 2 ) where X and C are the two input sequences and X ′ is the output of the Transformer layer . The Transformer layer in ( 2 ) adopts the original post-layer normalization architecture ( Vaswani et al. , 2017 ; Devlin et al. , 2019 ) that places layer normalization after residual connection , rather than pre-layer normalization ( Vaswani et al. , 2018 ; Wang et al. , 2019 ) . 3 Linear Unified Nested Attention ( Luna ) . Our goal is to design an efficient attention mechanism to solve the quadratic complexity problem of full attention . We first introduce the proposed linear unified nested attention mechanism , named Luna attention ( §3.1 ) , and the architecture of each Luna layer ( §3.2 ) . Then , we present the variant of Luna for causal attention , named Luna causal attention ( §3.3 ) . Finally , we discuss the differences between Luna and three closely related models : Linformer ( Wang et al. , 2019 ) , Set Transformer ( Lee et al. , 2019 ) ( §3.4 ) and Shared Workspace ( Goyal et al. , 2021 ) . 3.1 Pack and Unpack Attention . The key idea behind Luna is to decouple the regular attention function in ( 1 ) into two nested attention operations , both of which have linear efficiency . To achieve this , besides the original query and context input sequences , Luna introduces an extra input that is a sequence with fixed ( constant ) length . With this extra input as the query sequence , Luna uses its first attention , named pack attention , to pack the context sequence into a fixed-length sequence . Formally , let P ∈ Rl×d denote the extra input sequence with fixed length l. The pack attention first packs C to YP with P as the query sequence : YP = Attn ( P , C ) ( 3 ) where Attn ( · , · ) is the regular attention function in ( 1 ) , C ∈ Rm×d is the context sequence , and YP ∈ Rl×d is the output of the pack attention , which is named the packed context . Since the length of P is a constant l , the complexity of pack attention is O ( lm ) , which is linear with respect to m. To unpack the sequence back to the length of the original query sequence X , Luna leverages its second attention , named unpack attention : YX = Attn ( X , YP ) ( 4 ) where X ∈ Rn×d is the original query sequence . Similar to pack attention , the complexity of unpack attention is O ( ln ) , which is also linear with repect to n. Encoding Contextual Information in P . The next question is where the extra input sequence P comes from . One straightforward choice is to format P as a learnable parameter of each Luna layer . One obvious drawback of this method , however , is that P would not capture any contextual information . To enhance the capacity of the Luna model , we propose to formulate YP as an additional output of each Luna layer , corresponding to P . Formally , the Luna attention function LunaAttn ( · , · , · ) takes three sequences as input and generates two sequence as output : YX , YP = LunaAttn ( X , P , C ) ( 5 ) where the computation of YP and YX is in ( 3 ) and ( 4 ) . By stacking multiple layers of Luna attention , the output YP from the previous layer , which captures contextual information of C , is employed as the input P of the next layer . For the first layer of Luna , we formulate P as learnable positional embeddings2 ( Vaswani et al. , 2017 ) . Reducing the Number of Parameters . Due to the two nested attention operations , there are two sets of parameters ( WQ , WK , WV ) in a single Luna attention function . There are several techniques to reduce the number of parameters , such as parameter sharing ( Xia et al. , 2019 ) . In this work , we follow Wang et al . ( 2020 ) to share WK and WQ in each layer , and conduct experiments to analyze performance decline against Luna with full sets of parameters ( §4.2 ) . 3.2 Luna Layers . The Luna attention is used as a drop-in-replacement for the regular attention . We incorporate the position-wise feed-forward network and layer normalization into Luna layers . Concretely , layer normalization is applied to both YX and YP , while FFN only to YX : YX , YP = LunaAttn ( X , P , C ) XA , PA = LayerNorm ( YX +X ) , LayerNorm ( YP + P ) X ′ , P ′ = LayerNorm ( FFN ( XA ) +XA ) , PA ( 6 ) where X ′ and P ′ are the two outputs of the Luna layer . The graphical specification of one Luna layer is illustrated in Figure 2 . 3.3 Luna Causal Attention . As discussed in Tay et al . ( 2020b ) , the ability to support causal autoregressive decoding , i.e . attending solely to the past and current tokens , is required when designing efficient self-attention mechanisms . However , due to the pack attention that packs the long sequence X into a fixed ( shorter ) length , it is not straight-forward to support causal attention in Luna . To design causal attention in Luna , we need to assume that the input P contains no information of X , i.e . P will not leak any future information of X to the history . Before we describe the Luna causal attention mechanism , we first define a causal function f : Rn×d1 × Rn×d1 × Rn×d2 → Rn×d2 : F , f ( X , Y , Z ) , where Ft = 1 t Xt t∑ j=1 Y Tj Zj ( 7 ) where F ∈ Rn×d2 and Ft denotes the t-th row of F . From the definition of f in ( 7 ) , we see that Ft can only access the information of the past and present row of X , Y and Z . To perform Luna causal attention , we first compute the attention matrix of the pack attention : Apack = ω ( PX T / √ d ) . For simplicity , we omit the learnable parameters , e.g . WQ , WK , WV in ( 1 ) . Note that for Apack , we can not use the softmax function for ω , as the normalization term in softmax leaks future information of X to the history . Inspired by the causal attention mechanism in Linear Transformer ( Katharopoulos et al. , 2020 ) , we use two activation functions : 1 ) ω ( · ) = elu ( · ) +1 based on the exponential linear unit ( Clevert et al. , 2016 ) ; 2 ) ω ( · ) = softplus ( · ) based on the softplus function ( Glorot et al. , 2011 ) . With the causal function f in ( 7 ) , we compute the attention matrix of the unpack attention : Aunpack = ω ( f ( X , X , ATpack ) ) . Unlike Apack , we can use ω ( · ) = softmax ( · ) for Aunpack , because the normalization is along the l-dimension rather than the n-dimension of X . Finally , the output Y is computed by Y = f ( Aunpack , ATpack , X ) . The complexity of the causal attention in Luna is still linear : O ( ln ) . One drawback of Luna causal attention , similar to the causal attention in Random Feature Attention ( RFA ) ( Peng et al. , 2021 ) and Linear Transformer ( Katharopoulos et al. , 2020 ) , is its sequential computation for each timestep t. The sources of P . In the formulation of causal attention , P is expected to contain no information about X . Thus , we need to formulate P based on the usage mode of the causal attention . For the encoder-decoder mode in sequence-to-sequence modeling ( e.g . for machine translation ) , we can use packed output from the Luna encoder as P . For the decoder-only mode ( e.g . for language modeling ) , P might be formulated as a learnable parameter of each layer . 2We also experimented with sinusoidal positional embeddings , and obtained similar results .

Summary: This paper proposes Luna, an approximation to the Transformer's softmax attention with linear (in sequence length) time and space complexity. The approximation works using two attention functions. The first packs the sequence into one of fixed length; the second unpacks this fixed length sequence into one whose length matches the original sequence. Experiments show that Luna is computationally efficient, and typically performs either competitively or better than baseline methods.