Luna: Linear Unified Nested Attention
Authors: Xuezhe Ma
Xiang Kong
Sinong Wang
Chunting Zhou
Jonathan May
Hao Ma
Luke Zettlemoyer
Venue: NeurIPS 2021 Poster (2021)
Keywords: Efficient Attention
Transformer
Linear Attention

Abstract: The quadratic computational and memory complexities of the Transformer's attention mechanism have limited its scalability for modeling long sequences.  In this paper, we propose Luna, a linear unified nested attention mechanism that approximates softmax attention with two nested linear attention functions, yielding only linear (as opposed to quadratic) time and space complexity. 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. As compared to a more traditional attention mechanism, Luna introduces an additional sequence with a fixed length as input and an additional corresponding output, which allows Luna to perform attention operation linearly, while also storing adequate contextual information. We perform extensive evaluations on three benchmarks of sequence modeling tasks: long-context sequence modelling, neural machine translation and masked language modeling for large-scale pretraining. Competitive or even better experimental results demonstrate both the effectiveness and efficiency of Luna compared to a variety of strong baseline methods including the full-rank attention and other efficient sparse and dense attention methods. 

Luna: Linear Unified Nested Attention
Xuezhe Ma¹ ISI, USC xuezhema@isi.cm edu
Xiang Kong" LTT, CMU xiangkêcs. cmu. cmu.edu
Sinong Wang" Facebook Al
sinongsangêfb.com com
Chunting Zhou LTI, CMU
Jonathan May ISI, USC
Hao Ma, Luke Zettlemoyer Facebook (haom, 1sz} com
chuntinz@cs, cmu. edu
j jonmay@isi.c edu
Abstract
The quadratic computational and memory complexities of the Transformer's at- tention mechanism have limited its scalability for modeling long sequences. In this paper, we propose Luma, a linear unified nested attention mechanism that approximates softmax attention with two nested linear attention functions, yielding only linear (as opposed to quadratic) time and space complexity. As compared to more traditional attention mechanism, Luna introduces additional sequence with a fixed length as input and an additional corresponding output, which allows Luna perform attention operation linearly, while also storing adequate contextual information. We perform extensive evaluations on three benchmarks of sequence modeling tasks: long-contex: sequence modeling, neural machine translation and masked language modeling for large-scale pretraining. Competitive or even better experimental results demonstrate both the effectiveness and efficiency of Luna compared to variety of strong baseline methods including the full-rank attention and other efficient sparse and dense attention methods. The implementation of our model is available at https://github..cmm wwt.aaooo//oeeeeeeaaaaaaaaaaa
Introduction
256 de Luma-1
Luna 16
Transformers (Vaswani et al., 2017) are surprisingly versatile models that preform well wide range of language and vision tasks, including machine translation (Vaswani et al., 2017: Ott et al., 2018), language understanding (Devlin et al., 2019), im- age recognition (Dosovitskiy al., 2020) and bioin- formatics (Madani et al., 2020). Attention (Bah- danau et al., 2015) provides the key mechanism that captures contextual information from the entire se- quence 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 mech- anism w.r.t the length of the input sequence, which prohibitively restricts their potential application to tasks requiring longer input sequences.
Avg. LRA Score (wla Retrieval)
53
Blotlir
53
Syithesiier
51
Sinkham
Linformer 6 Performer
Reformer
Transformer
49
Local Attention
Relative Speed Comparision
Figure 1: Trade-off between accuracy (y-axis), speed (x-axis) and memory (cir-radius) on LRA.
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 al., 2019), strided attention patterns (Child et al., 2019; Beltagy al., 2020),
Equual contribution.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

compressed attention (Liu et al., 2018), and attention with learnable patterss (Kitaev et al., 2020; Tay al., 2020a: Roy et al., 2021). Another emerging upproach 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-dimensionl representation by assuming low-rank structure in the full-rank attention matrix. Recently, some kernel-base 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 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 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 fixed length as input and an additional corresponding output. Importantly, the extra input allows Luna to perform attention operation linearly 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 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 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 = AAttn(X,)) - (Wo((WW)))
CWv
(I)
Vd
where the attention function Attn : Raxd x R.mxd_ - Bend takes as inputs two sequences: the query sequence X E Raad with length and the context sequence C' E gmnd with length m. and output one sequence Y E guxd with the same length n as the query X. d is the embedding dimension, and Wo, WK, Wy E Idddd are three learnable parameters that project the input sequences into the space of query, key and value matrices: Q = XWQ, K = CWx, V = CWy. w is an activation function, e.g. the softmax function in regular attention. Note that the formulation in (1) is applicable to both pross-attention where C' and X 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), 7),, which performs the attention function h times in parallel, is commonly used. Throughout this paper, we omit for simplicity.
In particular, the matrix A = w( QK")
Vdz E RRax 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 takes (mm) time and space, which is quadratic with respect to the sequence length and becomes 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-wis wise feed- forward 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 X"
LayerNorm (ttn (X, C) + +X) LayerNorm( (XF((X)) +A)
(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).
2

---

X² e ER
Add & Norm
18/
Feed Forward Layer
Add & Norm
Add Norm
Add Norm
Layer
Multi-
Add & Norm
Q
Multi-head Amention
Multi-head Attention
f
V
e R
P 2
(a) Regular
(b) Unified (Luna)
Figure 2: Illustration of the architecture of one Transformer encoder layer (left) versus one Luna encoder layer (right).
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 al., 2019) ($3.4) and Shared Workspace (Goyal 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 a sequence with fixed (constant) length. With this extra input as the query sequence, Luna uses its first attention, named pack attention, pack the context sequence into fixed-length sequence. Formally, let P e g denote the extra input sequence with fixed length The pack attention first packs C Yp with the query sequence:
Yp Attn(P,
(3)
where Attn(:,-) is the regular attention function in (1), C Rmxd is the context sequence, and Yp R is the output of the pack attention, which is named the packed context. Since the length of a constant 1, the complexity of pack attention is (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, (X, Yp)
where X Raxd is the original query sequence. Similar to pack attention, the complexity of unpack attention is (In), 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 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 additional output of each Luna layer, corresponding P. Formally, the Luna attention function LunaAtin(,.>)) takes three sequences as input and generates two sequence as output:
Yx, Yp = LunaAttn(X,P,oO
(5)
where the computation 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
3

---

the input P of the next layer. For the first layer of Luna, we formulate P as learnable positional embeddings (Vaswani et al., 2017).
Reducing the Number of Parameters. Due the two nested attention operations, there are two sets of parameters (Wo. WK. Wv) in single Luna attention function. There are several techniques to reduce the number of parameters, such as parameter sharing (Xia et al., 2019). this work, we follow Wang et al. (2020) to share WK and Wo 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 drop-in-replasee repppin-rrlacmmentt for the regular attention. We incorporate the position- wise seed-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.P C)C)
XA. PA = LayerNorm (Yx + X). LayerNorm(Ypp + P) X P' LayerNorm(PFN(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.c. 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, is not werighh-ffrwwrt 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 causal function f: BRaxd x Raxxd RRxxd2 -> Ruxxd;
F f(X,Y,Z), where F. = X
D Y
t
j=1
where F' E [Ruxda and FÚ denotes the t-th row of F. From the definition of f in (7), we see that Fe 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 = w(PXT/ /Vd/ For simplicity, we omit the learnable parameters, c.g. WQ. WK. Wy in (1). Note that for Apack- we cannot use the softmax function for softmax leaks future information of to the history. Inspired by the causal attention mechanism in Linear Transformer (Katharopoulos et al., 2020), based on the exponential linear unit (Clevert et al., 2016); 2) w-)) = softplus(.) based on the softplus function (Glorot al., 2011). With the causal function. f in (7), we compute the attention matrix of the unpack attention: Aumpack = w( w([(X, X, Apack)) Unlike Apack, we can use w--) = softmax(-) for Aunpack. because the normalization along the l-dimension rather than the n-dimension of X. Finally, the output Y is computed by = [Auupackk AbackkX] X).
the normalization term in
use two activation functions: 1) w)- clu(-) +1
The complexity of the causal attention in Luna is still linear: (In). One drawback of Luna causal attention, similar to the causal attention in Random Feature Attention (RFA) (Peng et al., 2021) and Linear Transformer (Katharopoulos al., 2020), its sequential computation for each timestep
The sources of P. In the formulation of causal attention, 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-sequees 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.
modeling (e.g. for machine translation), we can use
We also experimented with sinusoidal positional embeddings. and obtained similar results.

---

Table Experimental results on the long range arena (LRA) benchmark. For Luna, we explore three projected dimensions: 16, 128 and 256. 'Avg. (w/o rtl) denotes the averaged accuracy over all tasks excluding Retrieval. The performance of previous works are from Tay al. (2021).
Models
ListOps
Text
Retrieval
Image
Pathfinder
Avg.
Avg. (w/o rtl)
Transformer
36.37
64.27
57 76
42.44
71.40
54,39
53.62
Transformer (re-ing (re-impl)
37.11
65.21
79.14
42.94
71.83
59.24
54.27
Local Attention
15.82
52.98
53.39
41.46
66.63
46.06
44.22
Sparse Trans.
17.07
63.58
59.59
44.24
71.71
51.24
49.15
Longformer
35.63
62.85
56.89
42.22
69.71
53.46
52.60
Linformer
35.70
53.94
52.27
38.56
76.34
51.36
51.14
Reformer
37.27
56.10
53.40
38.07
68,50
50.67
49.99
Sinkhorn Trans.
33.67
61.20
53.83
41.23
67.45
$1.39
50.89
Synthesizer
36.99
61.68
54.67
41.61
69.45
52.88
52.43
BigBird
36.05
64.02
59.29
40.83
74.87
55.01
53.94
Linear Trans.
16.13
65.90
53.09
42.34
75,30
50.55
49.92
Performer
18.01
65,40
53.82
42.77
77.05
51,41
50.81
Luna-16
37.43
65.74
79.38
46.39
78.36
61.46
56.98
Luna-122 128
38.01
65.74
79.55
47.47
78.89
61.93
57.53
Luna-256
37.98
65.78
79.56
47.86
78.55
61.95
57.54
3.4
Discussion
Relation Linformer and Shared Workspace. One previous work closely related Luna is Linformer (Wang et al.. 2019). Linformer linearly projects the context sequence C' E Rnxd into sequence with fixed length C' FC, where C" Rixd is the projected context sequence and E E R is the learnable projection matrix of each layer. Then, the attention operation is applied on the query X and the projected context C". The pack attention Luna a generalization of the linear projection in Linformer. There are two main advantages to Luna over Linformer: i) with pack attention as the projection method, Luna able to model sequences with various lengths. In contrast, Linformer requires the length of all input sequences to be the same due to the projection matrix E, whose shape depends on m. ii) Luna achieves better expressiveness than Linear, not only due to the general projection method but also by encoding adequate contextual information into the projection via P (see $3.1). Experimental improvements over non-contextual projection demonstrate the effectiveness of Luna (see $4.2). In contemporancous and individual work, Goyal et al. (2021) formulate contextual as a shared global workspace, which shares similar instantiation with Luna.
Relation Set Transformer. The additional input Luna can be regarded as a side memory module that can access the entire sequence to gather contextual information. From this view of point, Luna is also closely related to Set Transformer (Lee et al., 2019), an early model to integrate side memory module in Transformers. Similar to the projection matrix in Linformer, the inducing points in Set Transformer are learnable parameters. Thus, these inducing points might be formulated as the con-etaextual version of P in Luna. Moreover, Set Transformer is designed for set-input problems, which are problems wherein the input is a set of features and the model thereby invariant to permutation ordering of the input features (Tay et al., 2020b), while Luna attention is used as drop-in replacement for regular softmax attention.
Experiments
4.1 Long-Context Sequence Modeling
We evaluate the effectiveness and efficiency of Luna on the Long Range Arena (LRA) benchmark recently introduced by Tay et al. (2021), which designed for the purpose of evaluating efficient Transformer models under the long-context scenario. They collect five tasks in this benchmark which are ListOps (Nangia and Bowman, 2018), byte-level text classification (Text; Maas et al., 2011 1), byte-level document retrieval (Retrieval: Radev et al., 2013), image classification on sequences of pixels (Image; Krizhevsky et al., 2009) and Pathfinder (Linsley et al., 2018). These tasks consist of input sequences ranging from IK to 8K tokens and span across variety of data types and modalities.
5

---

Table 2: Training speed and peak memory consumption comparison of different models on byte-level text classification with various input lengths (IK, 2K, 3K and 4K). The best model is in boldface.
Steps per second
t
Peak Memory
Usage
(GB)
Model
1K
2K
3K
4K
IK
2K
3K
4K
Transformer
1.0
1.0
1.0
1.0
1.00
1.00
1.00
1.00
Local Attention
1.1
1.7
3.2
5.3
0.49
0.29
0.19
0.14
Linformer
1.2
1.9
3.7
5.5
0.44
0.21
0.18
0.10
Reformer
0.5
0.4
0.7
0.8
0.56
0.37
0.28
0.24
Sinkhorn Trans
1.1
1.6
2.9
3.8
0.55
0.31
0.21
0.16
Synthesizer
1.1
1.2
2.9
1.4
0.76
0.75
0.74
0.74
BigBird
0.9
0.8
1.2
1.1
0.91
0.56
0.40
0.30
Linear Trans.
1.1
1.9
3.7
5.6
0.44
0.22
0.15
0.11
Performer
1.2
1.9
3,8
5.7
0.44
0.22
0.15
0.11
Luna-16
1.2
1.8
3.7
5.5
0.44
0.23
0.17
0.10
Luna-122
1.1
1.7
3.4
5.1
0.49
0.28
0.21
0.14
Luna-256
1.1
1.7
3.3
4.9
0.60
0.33
0.23
0.16
To ensure fair comparisons, for all tasks except for the task Retrieval, closely follow the model configurations in Tay et al. (2021) such as data preprocessing, data split, model architecture, etc. For the task of Retrieval, we find that models are not fully converged when being trained for 5K steps as stated in Tay et al. (2021). Therefore, we train models for 20K steps for this task and obtain much better results. For a direct comparison, besides the average performance of models across all tasks, we also report the average accuracy on tasks excluding Retrieval. We run each experiment for five times with different random seeds and report the average accuracy. The syper-parameteer for each task are shown in Appendix A.1.
Results. The results of various models on the LRA benchmark are presented in Table 1. For our proposed method, we report results from models of three different projected dimensions (16, 128 and 256). First, we note that Luna achieves good results on all tasks consistently compared to the Transformer model and significantly outperforms all the other baseline methods in terms of the average accuracy. By taking a closer look at the accuracy for each individual task, Luna wins over baseline models on three out of five tasks and performs comparably with the best performed model on the other two tasks, i.e. ListOps and byte-level text classification. Notably, Luna improves over the Transformer model on image classification and pathfinder by a large margin. Second, we observe that although Luna achieves the best average performance with a projection dimension of 256, it also performs considerably well with smaller projection dimensions (16 and 128). This demonstrates the effectiveness of Luna even with small projected dimensions.
Memory and Speed Efficiency. Luna employs two nested linear attention functions to reduce the time and memory complexity compared to the vanilla softmax attention. Here, we examine the speed and memory footprint of various models with varying input lengths (1K, 2K, 3K and 4K). Following Tay al. (2021), all models are evaluated on the byte-level classification task with the same batch size. The result is shown in Table 2.
Considering the memory efficiency, Luna with projected dimension of 16 is highly memory- efficient, which is only 10% of the vanilla Transformer at 4K input sequence length. With larger projected dimensions, i.e 128 and 256, Luna requires more memory but is still competitive compared to other efficient Transformer models. In terms of time efficiency, Luna-16 speeds up over the standard Transformer by 1.2-5.5 times, varying by the sequence length. Compared other efficient Transformers, Luna- 16 performs comparably with the fastest models, i.e. Performer and Linformer. Overall, our models achieve competitive advantage both in time- and memory-eftie efficiency over other models, while attaining the best performance on the LRA benchmark (see Figure 1).
In addition, we plot the trade-off among memory, time and averaged LRA score without task Retrieval in Figure 1. Models such as Linformer and Performer have faster speed and small memory requirement with the sacrifice of performance. However, besides competitive time- and memory-efficieng Luna models retain superior performance even with a small projected dimension (!=16).

---

Contextual information in P of Luna. Recently, a popular method to model the classification task using Transformer- based models is to prepend a special sym- bol, [CLS]], to every input example. The last hidden state of this symbol is re- garded as the aggregate sequence repre- sentation. In Luna, we introduce an extra model input P which not only allows us to efficiently compute the attention mech- anism but learn contextual information as well. Theoretically, the P from the last layer is capable of learning the representation of the input sequence. To validate this, we extract P at the last layer and employ the mean pooling strategy over positions to obtain the final feature for classification. We test its performance i.e. ListOps, Text and Retrieval and report results in Table 3. We find that P-based methods obtain better scores across all tasks against the [CLS]-based one, validating the powerful ability of P to encode contextual information of the input sequence.
Table Performance comparison of two sentence repre- sentation methods on LRA benchmark.
Models
ListOps
Text
Retrieval
Avg.
Luna- 16, [CLS]
37.43
65.74
79.38
60.85
Luna-16, P
38.06
65.81
80.22
61.36
Luna-128, [CL.S]
38.01
65.74
79.55
61.10
Luna-128, P
38.27
65.89
80.27
61.48
Luna-256, [CLS]
37.98
65.78
79.56
61.11
Luna-256, P
38.36
66.07
80.25
61.56
three long-text modeling tasks in LRA (Tay et al., 2021),
4.2 Machine translation
To evaluate Luna on sequence-to-sequennc mod- eling. we conduct experiments on a standard machine translation benchmark, i.e. WMT' 14 English-Gennnn German (EN+DE) dataset (4.5M sen- tence pairs) The data split and preprocessing steps follow those of Vaswani et al. (2017), using the scripts from FairSeq (Ott et al., 2019). We share the source and target vo- cabularies within the language pair, with 37K byte pair encoding (BPE) types (Sennrich et al., 2016). The Luna models closely follow the architecture of Transformer base: encoder and decoder layers with 8 attention heads and dmodet/ dnidden = 512/2048. We train the Transformer-base model with two optimization methods: Adam (Kingma and Ba, 2015) and Apollo (Ma, 2020), and find Apollo achieves better performance. Therefore, we use Apollo the optimizer for all Luna models. For each experiment, we conduct distributed training across eight NVIDIA Tesla V100 GPUs with maximum batch size of 8192 tokens per GPU. Further details are provided in Appendix A.2.
Table 4: Test BLEU on WMT" DE.
Model
BLEU
# Param.
Transformer- (Adam)
27.8
64.9M
Transformer- (Apollo)
28.3
64.9M
REA (k- 256)
27.2
66.2M
Luna-16, elu, tied kv
27.1
69.6M
Luna-32, elu, tied kv
27.3
69.7M
Luna-16, soft plus, tied kv
27.3
69.6M
Luna-32, softplus, tied kv
27.5
69.7M
Luna-16 16, elu
27.4
77.5M
Luna-32, elu Luna-16, softplus Luna-32, softplus
27.6
77.6M
27.6
77.5M
27.8
77.6M
Results. Table 4 presents the results of Luna on the test set BLEU scores of WMT' 14 EN-DE, along with Transformer-bar and Random Feature Attention (RFA) as baselines. Different from Peng et al. (2021) where the random feature attention is applied only to decoders, the RFA model in Table 4 applies random feature attention in both the encoder and decoder for a fair comparison. k 256 is the number of feature maps in RFA. For Luna, we report performance of models with different 16 and = 32, different activation functions in (7): elu(-) + and softplus(-), and w./w.o parameter sharing.
projected lengths:
From Table 4, the first observation is that softplus(-) consistently outperforms elu(-) + 1. Thus, we use softplus(-) :) as the default activation function in the implementation. Another interesting observation is that Luna with small projected length (1 = 16) obtains similar performance RFA with k 256 feature maps. Luna with 1 = 32 achieves competitive performance, but still falls behind the Transfonmer-bass model. Further improving the machine translation performance of Luna is left to future work. We also report the number of parameters of different models. At last, we evaluate Luna w./w.o parameter sharing. Although there are two sets of parameters in a single Luna attention function (Wq. WK. Wv), as mentioned in $3.1, we tie Wt with W, to reduce the number of parameters, and the performance decline is marginal. As a result, Luna with shared parameters has 7% and 5% more parameters compared to the vanilla Transformer and RFA models.
7

---

Effect of Encoding Contextual Information into P. As dis- cussed in $3.4, one advantage of Luna against Linformer is to
Table 5: Dev and Test BLEU
incorporate contextual P by formulating it as extra input. To
Model
Dev.
Test
investigate the importance of this design, we conduct experiments on WMT" 14 to compare Luna with the baseline model where P is formulated as a non-contextual learnable parameter of each layer.
Non-Contextual
24.4
25.2
Contextual
25.9
27.3
For both the contextual and non-contextual models, we train Luna with ( = 16, parameter sharing and plus, Table 5 lists the BLEU scores on the development and test sets. Luna with contextual P significantly outperforms the baseline with non-contextual P, demonstrating the effectiveness of this design in Luna.
4.3 Masked Language Modeling for Large-Scale Pretraining
One popular application of Transformer is to pretrain a large-scale language model on large amount of data which can then be fine-tuned on a wide range of downstream tasks, such as BERT (Devlin et al., 2019), RoBERTa (Liu et al., 2019), etc. Therefore, we pretrain Luna-based language model with RoBERRTaaeaee version with BookCorpus (Zhu et al., 2015) and English Wikipedia (totally 16GB), 2) RoBERTa version with BookCorpus, English Wikipedia, CC-News (Nagel, 2016), OpenWebText (Gokaslan and Cohen, 2019) and Stories (Trinh and Le, 2018) (totally 160GB). For Luna models, we set = 128. On the larger training corpus (160GB), we train models w./w.o parameter sharing. respectively. We compare models with 1-BERTa-base. BERT-base and Linformer which are trained on the same training data. Experimental details are provided in Appendix A.3.
model configuration on two versions of data as our pretraining set: 1) BERT
Finetuning Luna After obtaining the pretrained Luna-based language model, we finetune it on various natural language processing tasks, including sentiment classification (SST-2: Socher et al., 2013), natural language inference (QNLI: Rajpurkar et al., 2016), textual similarity (QQP: Chen et al., 2018, question answering (RACE (Lai et al., 2017) and CommonsenseQAA (CSQA; Talmor et al., 2019). For GLUE tasks. following Liu et al. (2019), we consider a limited hyperparameter sweep for each task, with batch sizes E {16, 32} and learning rate E {5e 6 Le 2e "), with a linear warmup for the first 6% of steps followed by a linear decay to Finetuning performed for 20 epochs with early stopping based on each task's 's evaluation metric on the dev set For QA tasks, we concatenate each candidate answer with the corresponding question and passage. We then encode every candidate and pass the [CLS] output at the last layer through fully-connected layer, which is used to predict the correct answer. We truncate question-answe
pairs that are longer than 128 tokens and, if needed, the passage so that the total length is at most 512 tokens. Following Liu et al. (2019), we try a small range of possible values for hyperparameters, i..., batch size E {16, 32), learning rate E {le 2e 5, 3e 5) and dropout E {0.0, 0.1, 0.2}. For other configurations such as warm-up steps, optimizer, we follow thoses in Liu et al. (2019).
The result is reported in Table 6. We observe that on the smaller dataset (16GB) our Luna model has similar or slightly better downstream results compared to other pretrained language models. On QNLI and SST-2, Luna models obtain the best performance among all models, reaffirming the effectiveness of Luna in pre-training. This demonstrates the strong ability of Luna for language representations. On the larger dataset (160GB), however, the performance of Luna is slightly worse than RoBERTa with vanilla Transformer architecture. One possible reason is that the capacity of Luna is not as sufficient as vanilla Transformer, due to the efficient attention mechanism. This is supported by the evidence that Luna with full sets of parameters achieves better performance than that with parameter-sharie because Luna with full sets of parameters has better capacity.
5 Related Work
There has been signficiant prior work on improving the efficiency of Transformers, besides the three closely related works discussed in $3.4. The common techniques include, but are not limited to, weight sharing (Dehghani et al., 2018), quantization (Shen et al., 2020; Fan et al., 2020), sparse attention (Parmar al., 2018; Kitaev al., 2020), side memory module (Lee al., 2019; Gupta and Berant, 2020; Goyal et al., 2021), and low-rank or compressed context (Wang et al., 2019; Ainslie
We observed that Luna finetuning requires more epochs than vanilla Transformer (20 vs. 10). We also finetuned RoBERTa with 20 epochs but did not obtain better results.
8

---

Table Performance of various models on development set of benchmark natural language under- standing tasks. Bold face indicates best performance.
GLUE QNLI
QA
Model
data
SST-2
QQP
RACE
CSQA
BERT-base
16GB
92.7
88.4
89.6
64.2
53.3
RoBERTa-base
16GB
93.1
90.9
90.9
65.6
Linformer-122
16GB
92.4
90.4
90.2
Luna-128, tied kv
16GB
93.1
91.2
90.8
65.2
53.1
RoBERTa-base
160GB
94.8
92.8
91.9
73.50
63.61
Luna- 128, tied kv
160GB
94.3
91.5
91.2
71.50
61.48
Luna-128
160GB
94.6
92.2
91.3
72.25
62.08
al., 2020). In this section, we briefly review some recently proposed methods. For a detailed overview we refer the readers to Tay et al. (2020b).
Sparse Attention The general idea of these methods is that, instead of attending to the whole sequence, each token only access to a fixed, predefined range such as local neighborhoods and strided or "dilated" windows. Popular methods include local attention (Parmar al., 2018), blockwise attention (Qiu et al., 2019), strided attention patterns (Child et al., 2019; Beltagy et al., 2020), and compressed attention (Liu al., 2018). To make this range more flexible, Reformer (Kitaev al., 2020) employs a hash-based similarity measure to efficiently cluster tokens into chunks and Routing Transformer(Roy al., 2021) employ online k-means clustering the tokens. The Sinkhorn sorting Network (Tay et al., 2020a) exposes the sparsity in attention weights by learning to sort blocks of the input sequence.
Kernel Methods. recently popular method improve the efficiency of Transformers avoid explicitly computing the attention matrix A in (1) by re-writing with kemnels. Typical models leveraging kernelization are Linear Transformer (Katharopoulos et al., 2020), Performer (Choro- manski et al., 2020) and Random Feature Attention (Peng et al., 2021). Since kernels are a form of approximation of the attention matrix, they can also viewed as form of low-rank method (Choro- manski et al., 2020) that compresses the context shorter length, such as Linformer (Wang et al., 2019) and the proposed Luna model.
The simplest technique to reduce the complexity of Transformer is to chunk input sequences into fixed blocks, with the obvious disadvantage of losing contextual information from past chunks. As discussed in Tay et al. (2020b), these models can be regarded as fixed pattern models. Transformer-XI. (Dai et al., 2019) proposed natural extension to the blockwise method to connect these blocks via a recurrence mechanism. Compressive Transformer (Rae et al., 2020) further extends Transformer-Xl by maintaining fine-grained memory of past chunk activations, which are discarded in Transformer-XL. Technically, Luna can be adapted to a recurrence method, by simply using P as an inherent memory module to maintain the recurrence across segments.
Recurrence.
6 Conclusion
We have introduced Luna, a simple, efficient and effective linear attention mechanism used drop-in substitute for regular softmax attention. By introducing an extra input with the fixed length, Luna is capable of capturing adequate contextual information while performing attention operations linearly. On three sequence modeling tasks, i.e., long-context sequence modeling, neural machine translation, and large-scale pretraining and finetuning, Luna achieves comparable even better performance than variety of strong baselines, while acquiring prominent gains of efficiency in both speed and memory. In future work, we are interested in combining Luna with recurrence methods where P can be used running memory across segments of inputs. Another interesting direction would be to apply Luna to other tasks with long input sequences, such as document-level summarization and translation.

---

Acknowledgments and Disclosure of Funding
This material is based on research sponsored by Air Force Research Laboratory (AFRL) under agree- ment number FA8750- reprints for Government purposes notwithstanding any copyright notation therein. Xiang Kong was supported by U.S. DARPA AIDA Program No. 7A8850--18220001 contained herein are those of the authors and should not be interpreted as necessarily representing the official policies or endorsements, either expressed or implied, of Air Force Laboratory, DARPA or the U.S. Government.
The U.S. Government is authorized to reproduce and distribute
The views and conclusions
References
Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek. Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. Etc: Encoding long and structured inputs in transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 268-2844 284, 2020.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 1e.11660706550 2016.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In International Conference on Learning Representations (ICLR), 2015.
Iz Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05/500 2020.
Zihan Chen, Hongbo Zhang, Xiaoji Zhang, and Leqi Zhao. Quora question pairs. University of Waterloo, 2018
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:/904. 10509, 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Djork-Aréé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learn- ing by exponential linear units (elus). In International Conference on Learning Representations (ICLR), 2016.
Zihang Dai, Zhilin Yang, Yiming Yang. Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context. In Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics, pages 2978 7778-2988,11 2019.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations (ICLR), 2018.
Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre- Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171-4186 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:20 1929, 2020.
Angela Fan, Pierre Stock, Benjamin Graham, Edouard Grave, Remi Gribonval, Herve Jegou, and Armand Joulin. Training with quantization noise for extreme fixed-point compression. arXiv preprint XXXx:200000032200 2020.
10

---

Xavier Glorot, Antoine Bordes, and Yoshua Bengio. Deep sparse rectifier neural networks. In Proceedings of the fourteenth international conference on artificial intelligence and statisties, pages 315-323. JMLR Workshop and Conference Proceedings, 2011.
Aaron Gokaslan and Vanya Cohen. holOpen-ebfexrCCoppe 2019.
Openwebtext corpus.
URI: https://skylion000 githuu
Anirudh Goyal, Aniket Didolkar, Alex Lamb, Kartikeya Badola, Nan Rosemary Ke. Nasim Rahaman, Jonathan Binas, Charles Blundell, Michael Mozer, and Yoshua Bengio. Coordination among neural modules through a shared global workspace. arXiv preprint arXiv:200301119, 2021.
Ankit Gupta and Jonathan Berant. Gmat: Global memory augmentation for transformers. arXiv preprint arXiv:2006.03274, 2020.
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mnss: Fast autoregressive transformers with linear attention. In International Conference Machine Learning, pages 5156-5165. PMLR, 2020.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference Learning Representations, 2015.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04453 2020.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. Technical Report. University of Toronto, 2009.
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. Race: Large-scale reading comprehension dataset from examinations. Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, pages 785-794, 2017.
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set trans- former: framework for attention- based permutation-invane neural networks. In International Conference on Machine Learning, pages 3744-3753, PMLR, 2019.
Drew Linsley, Junkyung Kim, Vijay Veerabadran, Charles Windolf, and Thomas Serre. Learning long-range spatial dependencies with horizontal gated recurrent units. S. Bengio, H. Wal- lach, H. Larochelle, K. Grauman, N. Bisnch, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL http:://pocedddinggeerrrpr.ccppp
Peter Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. International Conference on Learning Representations (ICLR), 2018.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoeer and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprini 1rrXiv1900..11 1692, 2019.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, In International Confer- ence on Learning Representations, 2019.
Xuezhe Ma. Apollo: An adaptive parameter- wise diagonal quasi-newton method for nonconvex stochastic optimization. arXiv preprint arXiv:2009. J 13586, 2020.
Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies, pages 142-150, 2011.
Ali Madani, Bryan McCann, Nikhil Naik, Nitish Shirish Keskar, Namrata Anand, Raphael R Eguchi, Possu Huang, and Richard Socher. Progen: Language modeling for protein generation. bioRxiv, 2020.
Sebastian Nagel. Ce-news. URL: hup://web. archive. ongsuue/htp//commmerrrn..
org2016/00heevsaasscaaaiaabbbbb 2016.
11

---

Nikita Nangia and Samuel Bowman. Listops: A diagnostic dataset for latent tree learning. In Proceed- ings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Student Research Workshop, pages 92-99, 2018.
Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine translation. In Proceedings of the Third Conference on Machine Translation: Reseurch Papers, pages 1-9, 2018.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier. and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019: Demonstrations, 2019.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference Machine Learning, pages 4064. PMLR, 2018.
Hao Peng, Nikolaos Pappas, Dani Yogatama. Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. In International Conference on Learning Representations, 2021. URL http:://pponreeeeeeeee/fo/mmm/ QCTK/T/
Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self-attention for long document understanding. arXiv preprint arXiv:1911.02972. 2019.
Dragomir Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. The acl anthology network corpus. Language Resources and Evaluation, 47(4):919-944, 2013.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modeling. International Conference on Learning Representations (ICLR), 2020.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages
2016.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021.
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Proceedings of the 54th Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers). pages 1715-1725, 2016
Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami, Michael W Mahoney, and Kurt Keutzer. Q-bert: Hessian based ultra low precision quantization of bert. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 8815-8821, 2020.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang. Christopher D Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality In Proceedings of the 2013 conference on empirical methods in natural language processing, pages 1631-1642, 2013.
sentiment treebank.
Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2818-2826, 2016
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Commonsenseqa: A question answering challenge targeting commonsense knowledge. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 1111944455 2019.
Yi Tay. Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. In International Conference on Machine Learning, pages 1438-9447, PMLR, 2020a.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efffcciett transformers: A survey. arXiv preprint ar arXiv:2009.06733 06732, 2020b.
12

---

Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena benchmark for efficient transformers. In International Conference on Learning Representations, 2021. URL https: Vye
Tricu H Trinh and Quoc Le. A simple method for commonsense reasoning. arXiv preprint arXiv:1806. 028400 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Autention is all you need. In Advances in neural information processing systems, pages 5998-6008, 2017.
Ashish Vaswani, Samy Bengio, Eugene Brevdo, Francois Chollet, Aidan Gomez, Stephan Gouws, Llion Jones, Lukasz Kaiser, Nal Kalchbrenner, Niki Parmar, et al. Tensor2tensor for neural machine translation. Proceedings of the 13th Conference of the Association for Machine Translation in the Americas ( (Volume I: Research Track), pages 193-199, 2018.
Qiang Wang, Bei Li, Tong Xiao, Jingbo Zhu, Changliang Li, Derek F Wong, and Lidia S Chao. Learning deep transformer models for machine translation. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 1810-1822, 2019.
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04766 2020.
Yingee Xia, Tianyu He, Xu Tan, Fei Tian, Di He, and Tao Qin. Tied transformers: Neural machine translation with shared encoder and decoder. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 5466-5473, 2019.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization the transformer architecture. In International Conference on Machine Learning, pages 10524-10533. PMLR, 2020.
Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In Proceedings of the IEEE international conference computer vision, pages 19-27, 2015.
13


Reviews:
[{"review_id": "mCal7hjmGdS", "reviewer": "Reviewer_FUG1", "summary": "The authors propose a novel operation to replace full self attention by compressing the information of the input features to a smaller set of features using attention and then attending to this smaller set to generate the outputs. The authors provide experiments on synthetic benchmarks as well as real world NLP tasks showcasing the performance of the proposed attention operation and comparing it with various state-of-the-art efficient attention methods.", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n--------------\n\n- The idea of compressing the sequence using attention is simple and works very well\n- The performance on LRA is impressive both in term of accuracy and speed\n- The authors showcase competitive performance on commonly used datasets on machine translation and MLM\n\nWeaknesses\n-----------------\n\n- Although LRA is an interesting benchmark the performance comparison in terms of accuracy is relatively weak. All models seem to have high variance in performance and if standard deviation was reported the numbers would carry significantly less importance. That said, Luna performs significantly better than other methods (however it also performs significantly better than full attention (3 percentage points)). This is also observed in MLM and machine translation for which the performance is much closer to other efficient attention models.\n- Performance in terms of wall-clock time per epoch (or for the forward pass) is not provided for the real world datasets. For instance is Linformer faster in the MLM pretraining? Is RFA faster for machine translation? Also there is no performance comparison for the autoregressive Luna attention.\n\nI believe that the above are not significant issues, however they make it harder to place Luna attention in the performance/speed spectrum for real world datasets or autoregressive tasks. For instance, how much faster is Luna for machine translation given that it achieves 1 BLEU point lower performance than a traditional transformer.\n\n### Minor points\n\n- Although distinctly different, I believe that the idea is sufficiently related to Perceiver [1] such that it warrants a discussion let alone a mention.\n\n[1]: Perceiver: General Perception with Iterative Attention", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I would like a little more discussion on the limitations of autoregressive Luna attention. The authors, judging from their submitted code, do not seem to use the CUDA kernels of (Katharopoulos et al, 2020). Although that takes away from maintaining a custom CUDA kernel, what is the speed comparison for autoregressive tasks such as decoding for machine translation?\n\nNo societal impact discussion is needed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "bjcpRHadB-5", "reviewer": "Reviewer_Lezq", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This is a clearly written piece of work. The approach is novel and it is made clear how it differs from prior work. The method is a simple and elegant way to reduce both the time and space complexity of the Transformer's attention. The experiments show significant computational improvements (particularly for long sequences) over the naive attention mechanism, and are competitive against other efficient Transformer baselines. The empirical results are also strong, with Luna performing either competitively with or better than baselines on a wide range of tasks.\n\nHowever, the experiments section in particular leaves some room for improvement:\n- On several tasks, Luna outperforms the naive attention mechanism; this is surprising given that Luna is, in theory, an approximation to naive attention. It would therefore add value to the paper if some intuition behind these surprising results were provided.\n- The models in Table 2 are evaluated on the byte-level classification task, so presumably they do not use the causal attention architecture. However in lines 152-154, a limitation of the Luna causal attention is presented, namely that it requires sequential computation for each time step. It would therefore be useful to readers if the training speed were shown for the causal archiecture too.\n- In Section 4.2, certain hyperparameters for the pre-training procedure are not provided, e.g. batch size, learning rate.\n- The final part of Section 4.1 shows that using the final $Y_P$ and mean-pooling works better than using CLS, so why not use this strategy for the tasks in Section 4.2?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Besides the fact that Luna causal attention requires sequential computation for each time step, limitations of the method are not presented.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "KGiPS4mgEbv", "reviewer": "Reviewer_1DnM", "summary": "The paper proposes a linear time and memory attention mechanism that computes attention into two steps: In first step, an extra input ($P$) with fixed sequence length is introduced that summarizes the context to fixed length output sequence via packing attention. In the next step, the queries attend (unpack attention) to the fixed length summarized sequence resulting in a overall complexity that is linear in input sequences length. Experiments are conducted on LRA benchmark and other NLP tasks to show the effectiveness of the proposed method against other fast attention mechanisms. ", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- **Originality**:\n \n  - The paper presents a linear attention mechanism that first packs the context into a fixed length representation and then queries attend to the packed sequence to get the outputs. While for the self-attention, Set Transformers have been proposed using the same mechanism, the novelty of the paper lies in generalizing this mechanism to non-trivial causal auto-regressive modeling.  Moreover, as opposed to using learnable parameter to model the $P$ sequence for each layer, the paper also proposes to formulate $P$ as  learnable positional embedding for only first layer and the output of packed attention is fed as the input sequence $P$ to the next layer. This is later shown to have impact on the performance.\n\n- **Quality**:\n\n  - The paper is well organized, appropriately discusses the relation to other works, and discusses well different aspects of Luna attention such as the sources of $P$. I also appreciate the experiment validating the contextual information in $P$ as well as comparing the effect of encoding $P$ which is an important design parameter proposed in the work.  \n  - The experiments on LRA across other domains showcase the effectiveness of Luna against full attention and other fast attention mechanisms.  \n  - Experiments on the pre-training tasks and machine translation demonstrate the effectiveness on difficult large scale settings.\n\n- **Clarity**:\n  - The paper is well written and easy to follow. The paper discusses differences with closely related works such as Set Transformers and Linformer. Ablations analyzing contributions from different aspects of the work are also well analyzed and described. All details including the code should be enough to reproduce this work.\n\n- **Significance**:\n  - This paper tackles important problem of improving the computational efficiency of attention. The idea is well presented and generalizes the packing and unpacking attention to causal scenarios which is an important distinction with Set Transformers.  Experimental results on LRA show that compared to other fast attentions, Luna offers favorable trade-offs for time, memory, and performance benchmarks.\n\n - **Suggestions for improvement**:\n   - Calling method \"approximation\" of softmax attention seems a bit misleading.  Approximating softmax gives a connotation of being backward compatible or attention output approximating softmax attention patterns. In the absence of any empirical results or theoretical arguments towards this, I would suggest to rephrase this more appropriately.\n\n   - I would encourage the authors to include the following two relevant work in related works discussion:\n      - \"SMYRF: Efficient Attention using Asymmetric Clustering\" \n      - \"Fast Transformers with Clustered Attention\" \n\n   - Discussion on limitations:\n      - The paper currently is missing any discussion on potential limitations imposed by restricting the length of $P$. For instance, can the $P$ pose issues if a task requires to copy tokens from context and the number of tokens to be copied is more tokens than the $P$ length? Or would more layers can overcome this limitation? It would interesting to do some analysis on the sequence duplication synthetic task as defined in Reformer.\n\n    - Visualization of information captured by $P$.\n      - It would be interesting to see while packing to small context, how is attention distributed over the context tokens. A few visualization on the question answering tasks could be very insightful.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Not applicable", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8-9 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a linear attention based on a pack and unpack mechanism of an input sequence. Experiments on LRA, pretraining and finetuning, and machine translation demonstrate the benefit of the proposed approach. I think this is a useful addition to the linear attention literature. Al reviewers generally agree that this is a good paper and provided some suggestions, which the authors mentioned will be added to the paper. I recommend accepting this paper.
Decision: Accept (Poster)