cosFormer: Rethinking Softmax In Attention
Authors: Zhen Qin
Weixuan Sun
Hui Deng
Dongxu Li
Yunshen Wei
Baohong Lv
Junjie Yan
Lingpeng Kong
Yiran Zhong
Venue: ICLR 2022 Poster (2022)
Keywords: Linear Transformer
softmax attention

Abstract: Transformer has shown great successes in natural language processing, computer vision, and audio processing. As one of its core components, the softmax attention helps to capture long-range dependencies yet prohibits its scale-up due to the quadratic space and time complexity to the sequence length. Kernel methods are often adopted to reduce the complexity by approximating the softmax operator. Nevertheless, due to the approximation errors, their performances vary in different tasks/corpus and suffer crucial performance drops when compared with the vanilla softmax attention. In this paper, we propose a linear transformer called cosFormer that can achieve comparable or better accuracy to the vanilla transformer in both casual and cross attentions. cosFormer is based on two key properties of softmax attention: i). non-negativeness of the attention matrix; ii). a non-linear re-weighting scheme that can concentrate the distribution of the attention matrix. As its linear substitute, cosFormer fulfills these properties with a linear operator and a cosine-based distance re-weighting mechanism. Extensive experiments on language modeling and text understanding tasks demonstrate the effectiveness of our method. We further examine our method on long sequences and achieve state-of-the-art performance on the Long-Range Arena benchmark. The source code is available at https://github.com/OpenNLPLab/cosFormer.

Published as a conference paper at ICLR 2022
COSFORMER: RETHINKING SOFTMAX IN ATTENTION
Zhen Qin
1,3 Weixuan Sun
Hui Deng"
Dongxu Li Yunshen Wei BBaohong Ly
Juujjie Yan 2 Lingpeng Kong
1,2 Yiran Zhong"
"Sense Time Research
Shanghai Al Laboratory Australian National University
Northwestern Polytechnical University
The University of Hong Kong
flast f irrtittaatt tamttameeeaa.ommm r
com, wpkkcs..kkk hku.hk
ABSTRACT
Transformer has shown great successes in natural language processing, computer vision, and audio processing, As one of its core components, the softmax atten- tion helps to capture long-range dependencies yet prohibits its scale-up due to the quadratic space and time complexity to the sequence length. Kernel methods are often adopted reduce the complexity by approximating the softmax operator. Nevertheless, due to the approximation errors, their performances vary differ- ent tasks/corpus and suffer crucial performance drops when compared with the vanilla softmax attention. In this paper, we propose a linear transformer called COSFORMER that can achieve comparable or better accuracy to the vanilla trans- former in both casual and cross attentions. COSFORMER is based on two key properties of softmax attention: i). non-negativeness of the attention matrix; ii). non-linear re-weighting scheme that can concentrate the distribution of the at- tention matrix. As its linear substitute, COSFORMER fulfills these properties with linear operator and cosine-baset distance re- re-weighting mechanism. Extensive experiments on language modeling and text understanding tasks demonstrate the effectiveness of our method. We further examine method long sequences and achieve state-of-the-ee performance on the Long- Rangg Arena benchmark. The source code is available at COSFORMER
INTRODUCTION
56
BigBrdd
cosFormer
Longformer
Synthwsizer
Linformrr er
Long-Range
Sparse Transfomer.
Reformer
Perfomee
Linar Transfomar
Local attention
Speed (examples per sec)
Figure 1: Performance (y axis), speed (Ir axis), and memory footprint (circle sizes) of efficient transformers on the Long-Range Arena benchmark. The proposed CosFORMER achieves an all-around supremacy over competing methods in the top left quadrant.
With years of development, the transformer model (Vaswani et al. 2017 and its variants (Zaheer et al. 2020; Wang et al., 2020) Tay et al., 2020a) have been successfully adapted to three most popular artificial intelligence (Al) fields: Lee, natural language processing (Devlin
al. 2019
Liu et al. 2019), computer vision (Dosovitskiy et al., 2020; Carion et al. 2020 Liu et al. 2021
and audio processing (Schneider al. 2019 Baevski al. 2020). Compared with conventional
"nddcctess the corresponding author. t Indicates equal contribution.

---

Published as a conference paper at ICLR 2022
recurrent (Hochreiter & Schmidhuber 1997) and convolutional architectures (He et al. 2
transformer- based architectures are generally more scalable to data volumes (Brown et al., 2020) and stronger in capturing global information with less inductive bias, thus excelling many tasks.
Dot-product attention with softmax normalization is the cornerstone of the transformer to capture long-range dependencies. However, its quadratic space and time complexity with regard to the length of the sequence make its computational overhead prohibitive, especially for long inputs. To address this issue, numerous methods are proposed recently, such as the sparse attention matrix (Za-) heer et al. 2020; BBlltgy et al., 2020; Tay et al. 2020a; Kitaev et al. 2019) Child et al. 2019). low- rank representations (Wang et al. 2020) or bareel-based methods (Peng et al.) 2020, Choromanski et al. [020] Katharopoulos et al. 2020), among many others. These methods achieve reduced
computational complexity with comparable performances when compared with the vanilla attention
architecture on several selected tasks or corpus.
However, the improved efficiency is usually achieved via introducing additional yet often impracti- cal assumptions the attention matrix (Wang al. 2020) or with valid approximation of softmax operation only within constrained theoretical bounds (Choromanski et al. 2020; Peng et al. 2020)
Therefore, when their assumptions are unsatisfied or when approximation errors get accumu lated, these methods may not always be advantageous over the vamilla architecture (Narang et al. 2021). Consequently, performance deficiencies in broad application spectrum are often observed in these transformer variants, especially those with linear complexity. For example, the Performer (Choro- manski et al.) 2020). RFA (Peng al., 2020) and Reformer (Kitaev et al.; 2019) show less satis- factory performance on the GLUE benchmark (Wang et al. (0188) when compared with the vanilla architecture suggested in our preliminary experiments (Tab. Furthermore, many of these afore- mentioned methods are not applicable to casual attentions, which are critical for auto-regressive training. For example, techniques proposed in Linformer (Wang et al. 2020) and BigBird (Zaheer et al. 2020) are specific to cross attentions.
Since the softmax operator appears be the main hurdle while efficient yet accurate approximation to softmax is difficult to achieve, one question naturally arises: "Can we replace the softmax opera- for with a linear function instead, while maintaining its key properties?". By digging into the soft- max attention, we find two key properties that affect its empirical performance: (i) elements in the
attention matrix are non-negative (Tsai et al. 2019 Katharopoulos et al. 2020); (ii) the non-linear
re- weiggghttng scheme acts a stabilizer for the attention weights Tiississ 2016 Gao & Pavel 2017;
Jang et al. 2016). These findings reveal some new insights of the current approaches. For example, the linear transformer (Katharopoulos et al. 2020) achieves property (i) using an exponential linear unit (Clevert et al.) 2016) activation function. However, due to lack of the re-weighting scheme, it underperforms other efficient transformer variants on the Long-Range Arena benchmark as shown in Figure I as well as the language modeling task (Table based on our controlled experiments.
In this paper, we propose new variant of linear transformer called COsFORMER that satisfies both of the above properties. Specifically, we enforce the non-negative property by passing the features ReLU (Agarap) 2018) activation function before computing the similarity scores. In this way, we encourage the model to avoid aggregating negatively-correlatee contextual information. Further, we adopt cos re- weighting scheme to stabilize the attention weights. This helps the model to amplify local correlations, which usually contain more relevant information for natural language tasks. Thanks to the Ptolemy's 's theorem, our attention can be exactly decomposed into a linear form. We perform extensive experiments on both autoregressive language models and bidirectional models on five public benchmarks, including WikiText-103 (Merity et al. 2017). GLUE (Wang et al. 2018). IMDB Maas al. 2011). AMAZON (Ni et al. 2019) and Long- Range Arena benchmark (Tay et al. 2020b). Our model shows much better inference speed and smaller memory footprint, while achieving par performance with the vanilla transformer. It is noteworthy that our method ranks 1st on the Long-Range Arena benchmark, showing favorable performance than other competitors, which well demonstrates its strong capacity in modeling long sequence inputs.
2
OUR METHOD
In this section, we provide technique details of our linear transformer called COSFORMER. The key insight of the COSFORMER is to replace the non-decomposable non-linear softmax operation by a linear operation with decomposable non-linear re-weighting mechanism. Our model is applicable
2

---

Published as a conference paper at ICLR 2022
to both casual and cross attentions with linear time and space complexity with regard to the input sequence length, thus exhibiting strong capacity in modeling long-range dependencies.
2.1 THE GENERAL FORM OF TRANSFORMER
Given an input sequence æ with length of N, we first represent it in the embedding space x E Nxx with feature dimension of d. A transformer block T Nxd + RNxx xx with input x: is defined as:
= F(A(x) +
(1)
where F is feedforward network that contains a residual connection; A is the self-attention func- tion that computes the attention matrix A, which has quadratic space and time complexity with respect N, thus becoming the computation bottleneck of on long inputs.
There are three key components in A, namely, query (Q), key (K), value (V) computed through three learnable linear matrices WQ. K. Wv: Q zWQ. K xWK, , V zWv.. We use M, to represent the i-th row of a matrix M, then the output O E2 Nxd of A(x) can be computed as:
S(Q,, V,, 0.=5 Σ, Σ, S(Q,, K₁)
O = A(x) O ONI...,ON''", O, Σ
(2)
where S(-) measures the similarity between queries. If S(Q, K) oxp(QK"), the Eq. becomes the dot-product attention with softmax normalization. In this case, the space and time complexity to compute one row of the output O is O(N). Therefore, the total space and time complexity for computing O grows quadratically with respect the input length.
2.2 LINEARIZATION OF SELF-ATTENTION
According to Eq. 2 we can select any similarity functions to compute the attention matrix. In order to maintain linear computation budget, one solution is to adopt a decomposable similarity function such that:
S(Q,, K2) ((Q))o(
where & is a kernel function that maps the queries and keys to their hidden representations. one can rewrite Eq. in the form of kernel functions
Then
Σ (C Σx ((QQ)00K))
O
(4)
After that, attention operation in linear complexity is achieved via the matrix product property:
(o())o(K)*) = =())((()V))
(5)
In this form (Eq. 5). instead of explicitly computing the attention matrix A = QK NNN, we calculate the of o(K)'V E Rdddd first, and then multiplying o(Q) E Rxd By using this trick, we only incurs a computation complexity of O(Nd²). Note that in typical natural language tasks, the feature dimension of one head is always much smaller than the input sequence length N (d < N), so we can safely omit d and achieve computation complexity of O(N), as illustrated in Figure 2
Previous Solutions As aforementioned, the key to the linear attentions similarity function S(-) that generalizes well to different tasks. Most existing linear transformers are trying to find an unbiased estimation of the softmax attention. For example, RFA (Peng al., 2020) approximates the softmax operation with random feature maps using theorem of random fourier features (Rahimi & Recht] (2008) and the Performer (Choromanski et al., (020)) utilizes positive random features to approximate it. However, to the selection of sampling rate and becomes unstable if the sampling rate gets too high. Also, to accommodate recency bias, gating mechanisms are employed to better exploit more recent context.
find decomposable
empirically find that these methods are sensitive
Another group of works attempt to directly replace the softmax with a linear operation. For example, the linear transformer (Katharopoulos et al. 2020) model replaces the softmax similarity function with a pure dot product S = QK" and use a non-linear activation function o(-) = elu(-) + 1 to model the pairwise relation between features. However, solution does not necessarily generalize well on many downstream tasks (Tab. or the Long- Range Arena benchmark (Tab. 4. this paper, we propose a new replacement of softmax that not only achieves comparable or better performance than the softmax attention in a wide range of tasks, but also enjoys linear space and time complexity.
controlled experiments show that their
3

---

Published as a conference paper at ICLR 2022
O(Ned)
O(Nd?)
O(N'd)
O(NdF)
dkN
83
Nxd
NxN
8
Nxd
Nxd
X3
X3
Nxd
Nxd
Softmax
KT
Q'
OINd + Nod) O(N)
OINe NdF) O(N)
Vanilla self attention
Linearized self attention
Figure 2: Illustration of the computations for vanilla self attention (left) and linearized attention (right). The input length is N and feature dimension is d, with d < N. Tensors in the same box are associated for computation. The linearized formulation allows ((N) time and space complexity.
2.3 ANALYSIS OF SOFTMAX ATTENTION
In the vanilla transformer architecture, when S)Q exp(QKT) the softmax operation is ap- plied to obtain row-wise normalization on the attention matrix A E NN as shown in the Eq. 2 In other words, normalize the relations of each element in the input sequence all other elements in order to obtain weighted aggregation of contextual information. However, apart from the good empirical performance of softmax attention, what are the crucial and necessary characteristics of remain only loosely determined in the original transformer paper and follow-up works.
In this work, we empirically identify two key properties of the softmax operation that may play important roles for its performance: 1) ensures all values in the attention matrix A to be non-negative; 2) provides a non-linear re- weighting mechanism
Analysis the softmax properties. at- tention variants are implemented the RoBERTa (Liu) et al. 2019) architecture and are pre-trained on the Tikt-ext-103 (Merity al. 2017) dataset. The Loss represents the validation loss. We then fine-tune these variants on each downstream datasets and show the ac- curacy (the higher the better).
Table 1:
concentrates the dis- tribution of attention connections and stabilizes the training Titsias 2016 Gao & Pavel 2017
Loss
QQP
SST-2
MNLI
Jang et al. 2016).
2.343
84.23
76.26
58.27
To validate these assumptions, design the following preliminary studies as shown Ta- ble [ First, to validate the importance of non- negativity, we compare three instantiations of the function ó in equation 3 an identify mapping Φ1 that does not preserve the non-negativity, and the other variant (neLU(() = ReLU(-) that retains only positive input values while replacing negative values to zeros. We also add the (LeskyReLU()) = LeakyReLU()) variant as it does not well but have the same non-linearly the ReLU one. Second, demon-
&eeakyBLLU
2.246
84.46
78.21
74.26
QNLLU
1.993
88.86
89.90
77.86
softmax 1.915
88.41
92.31
79.15
have the non-negativvty strate the effect of non-linear re-weighting, we compare the models using only ΦReLU(.) without any re- weighgitting and those with softmax operations. From Table [ the superior results of ΦRELU over é1 and Re..skyRel.U demonstrate the benefit of retaining non-negative values. Our conjecture is that by retaining only positive values in the similarity matrices, the model ignores features with negative correlations, thus effectively avoiding aggregating irrelevant contextual information. By comparing the results of ΦRELU with the softmax, we observe that models with softmax re- re-weighting converge faster and generalize better to downstream tasks. This might be explained as softmax normalization amplifies the correlated pairs, which might be useful to identify useful patterns.
2.4 COSFORMER
Based on the observations above, we propose our model COSFORMER which discards entirely the softmax normalization while still features the non-negativity and re-weighting mechanism. Our COSFORMER consists two main components: a linear projection kernel Ólineur and a cos- cos-Based Re- weighting mechanism. Below describe details of each components:
Linear projection kernel Qlinaar Recall the general form of the attention in Eq. 2 let us define linear similarity as:
S(Q, K) = s(Cuneat (Q). Qlineur (K) = s(Q K')
(6)
where @lnnaar is the transformation function that map queries Q and keys K to desired represen- tations Q and K', and is function that can be linearly decomposed to measure the similarity between Q' and K'. Specificallly in order to ensure full positive attention matrix A and avoid
4

---

Published as a conference paper at ICLR 2022
555
(1)Vanilla
(2)cosFormer
[3)cosoormer (w/o re-weighting)
(4)cos re-weighting matrix
Transformer
Figure 3: (1): Attention matrix of vanilla (2n matrix of COSFORMER .(3): Attention
matrix of COSFORMER without re-weighting. (4): Visualization of the cos-based distance matrix. After weighting, see smoother attention distribution along the diagonal region of attention matrix, exhibiting a similar pattern to the vanilla transformer, which assists stabilize the training.
aggregating negatively-correlated information, we adopt ReLU(-) as the transformation functions and therefore effectively eliminate negative values:
dlinear (x) ReLU(x)
As Q' and K' contain only non-negative values, we directly take their dot-product s(x, v) my"" 3. Ixd followed by row-wise normalization compute attention matrices:
=
ΣN (chwoor (.tteer K
ΣY (RELU [dhnne (h (K2]) 25 (RLLU())))))
0
(8)
Based on Eq. we rearrange the order of dot-product and obtain the formulation of the proposed attention in linear complexity as:
ReLU(QU) EN ReLU(KK)* O = ReLUIQ.) E ReLU(K,)
cos-Based Re-weighting Mechanism The non-lineat re- reighttgtnng mechanism introduced by the softmax attention can concentrate the distribution of the attention weights and therefore stabilize the training process (Titsias] 2016 Gao Pavel 2017: 7; Jang et al. 2016). We also empirically find that
it can punish far-away connections and enforce locality in some cases. In fact, such locality bias, i.e,, large portion of contextual dependencies are from neighboring tokens, commonly observed on downstream NLP tasks (Clark et al. 2019; Kovaleva et al., 2019), as shown in Figure 3 (1).
Based on the assumption above, what we need to fulfill the second property of softmax may be a de- composable re-weighting mechanism that can introduce recency bias to the attention matrix. Here, we propose a cos-based re- re-weighting mechanism as it perfectly fit our purpose: 1). the Ptolemy's theorem ensures the cos weights
be decomposed into two summations; 2). shown in Fig- ure [ (4), the cos will put more weights on the neighbouring tokens and therefore enforces locality. Also, by comparing the attention matrices in Figure (2) and (3), the COSFORMER enforces more locality than the one without the re-weighting mechanism.
Specifically, by combining with Eq the model with cosine re-weighting is defined as:
s sQ. K cos (%
i-j
(10)
M
By leveraging the Ptolemy's theorem, we decompose this formulation as:
K,K,T M
a {-5 2
Q.K;T
cos
xi 2M
2M
+ sin
2M
2M
Q; cos
xi
K
coss
2dd
Q
2M
ar
(K)
sin
2M
where 1 1, ..., N, M > N, and Q" ReLU(Q), K" ReLU(K). Let Q Q' cos
Q = Q sin TM
mi
Koos = K cos
njj ZM
Ksin
K sin
ボj 237
the output the proposed at-
tention module can expressed as:
ΣN SQ ΣN x [(I,((K))))
DN Q: ((K)Te) V/z) + +E Q:"h ((KT"))
0,
(11)
EN (Kg00) +EN Φ: ":"((gi")""
5

---

Published as a conference paper at ICLR 2022
Vanila Tranfomerr
Van Tansforer
suxformer wio noowight
coeFormer reweggt
Valie Lave
22
22
13
Train Stop
Train Shop
Figure 4: Training loss (left) validation loss (right) of the bidirectional language modeling pre-train. In both training and validation, the proposed COSFORMER has a faster converge speed than vanilla transformer.
where O2 is the output at the tth position of the sequence from the attention module. Detailed derivation are included in the Appendix. Without losing the generality, our method achieves linear complexity as:
O = S(Q,K)V -Q006 Kooss + Qsin K V Q00 (Ko + Q (K
(12)
Relation to positional encoding. COSFORMER can be seen as a new way of introducing the rela- tive positional bias to the efficient transformer. Compared with the Rotary Position Embedding (Su et al. 2021), they use a more complex position embedding strategy and did not enforce the non- negativity to the similarity scores as ours. Also, since they only change the position embedding on the numerator while keeping the denominator unchanged, the summation of their attention scores is not equal to 1. For Stochastic Positional Encoding (Liutkus et al. (02T), I they use sampling strategy to approximate the softmax, and introduce relative positional encoding linear transformers.
3 EXPERIMENTS
In this section, we experimentally validate the effectiveness of the proposed method in multiple set- tings. The purposes of the experiments are three-fold. First, we validate the capacity of COSFORMER in language modeling through autoregressive (Sec. 3.1 and bidirectional (Sec. 3.2) setups using WikiTeext-103 103 (Merity et al., 2017). In this way, we validate the effectiveness of the proposed linear attention module in both causal and non-causal cases. Second, we investigate the generalization ability of COSFORMER on downstream tasks by comparisons with other existing transformer vari- ants. This is achieved by performing comparative finetuning experiments on five datasets, including GLUE (QQP, SST-2, MNLI) (Wang al., 2018). IMDB (Maas et al., 2011) and AMAZON (Ni) cet al. 2019) (Sec. 3.3). We further compare COSFORMER with other transformer variants on the lang-range-arena benchmark (Tay et al. (0206) to understand its ability in modeling long-range dependencies (Sec. (3.4) and show comparative analysis into model efficiency (Sec. 3.5) Third, we
conduct ablation studies understand each component in COSFORMER (Sec. 3.6
3.1 AUTOREGRESSIVEE LANGUAGE MODELING
In autoregressive or left-to-right language modeling. we estimate the probability distribution of token given its previous tokens. We use (Baevski Auli 2018) as our baseline model. Specifically, adopt their large model which has 16 cascaded layers with projected dimensions of 1024, and replace the self-attention module with our proposed linear attention module. We train our model on 8 Nvidia Tesla A100 GPUs with sequence length of 512 for 150K updates on the WikiText-103 (Merity et al.) 2017) and report perplexity on the validation and test splits in Table
We observe that although the baseline model a powerful standard transformer which re- quires quadratic computation complexity, cos- FORMER outperforms it with a clear margin in linear computation complexity. Besides, achieve comparable perplexity to other meth- ods on the validation set, and significantly out- perform all competing methods on the test set by clear gap, which further demonstrates the effectiveness of COSFORMER.
Table 2: Perplexity (lower is better) results of language modeling pre-training task on validation set and set the WikiText- 103 (Merity et al. (017)) dataset.
ppl(val) 1 ppl(test)
Vamilla Transformer
24.5
26.2
Linear Transformer
28.7
30.2
RFA- IPA-Gaussian
25.8
27.5
RFA-across
26.4
28.1
RFA-Gate-across
24.8
26.3
Trta-Gae--aausee
23.2
25.0
COSFORMER
23,5
23.1

---

Published as a conference paper at ICLR 2022
Table 3: Results on fine-tuned downstream tasks hased on pre-trained hidirectional model. Best result boldface and sccond best is underlined. The proposed COSFORMER achieves superb performances over competing efficient transformers and is approaching vanilla transformer.
QQP
SST-2 1
MNLI
IMDB 1
AMAZON 1
Avg 1
Vanilla Transformer (Liu et al., 2019
88.41
92.31
79.15
92.86
75.79
85.70
Performer (Choromanski el 2020
69.92
50.91
35.37
60.36
64,84
56.28
Reformer Kilaev ef 2019
63.18
50.92
35.47
50.01
64.28
52.77
Linear Trans, (Katharppollos et al.
2020
74.85
84.63
66.55
91.48
72.50
78.00
Longformer (Beitagy et al. 2020
85.51
88.65
77.22
91.14
73.34
83.17
RFA (Peng al. 2020
75.28
76.49
57.6
78.98
68.15
71.30
COSFORMER
89.26
91.05
76.70
92.95
76.30
85.25
3.2 BIDIRECTIONAL LANGUAGE MODEL.
For bidirectional language modeling, we adopt RoBERTa (Liu et al. 2019) as the baseline model. Similarly, we replace the self-attention module in the RoBERTa by the proposed linear attention module, and keep other structures unchanged. We train this bidirectional task on 2 Nvidia Tesla A100 GPUs for 50K iterations with input sequence length 512. As shown in Figure 4 COS- FORMER converges faster than vanilla transformer on both training and validation sets with com- parable or smaller loss values, despite it only consumes linear space and time computation com- plexity. In addition, the COSFORMER variant with re- re-weighting mechanism has both notably better converge speed and final results over the counterpart without re-weighting, which further validates the effectiveness of our cos-based distance matrix and also demonstrates the effectiveness of recency bias on natural language data.
3.3 DOWNSTREAM FINE-TUNING TASKS
In this section, we fine-tune the pre-trained model on downstream tasks demonstrate the general- ization ability of COSFORMER on downstream tasks. We use the pre-trained bidirectional model and fine-tune on three downstream text classification tasks: GLUE (QQP, SST-2, MNLi) (Wang al. 2018). IMDB (Maas et al., 2011) and AMAZON (Ni et al. 2019). For fair comparison, we first pre-train all the competing efficient transformer variants for the same 50K iterations on WikiText- 103 (Merity et al. (017)) under the same setting, then we follow the same fine-tuning protocol as RoBERTa (Liu et al. 2019) to fine-tune these methods on the downstream tasks. From Table we see that COSFORMER outperforms baseline (Liu et al. 2019) three out of five datasets, and achieves either best or secondary place on all five downstream datasets compared to competing effi- cient transformers. It is worth noting that despite Longformer (Beltagy et al.) 2020) achieves better results on MNLI than COSFORMER it requires a computation complexity of O(Nw) w), where w is window size. As shown in Figure Longformer is slower and requires more memory overhead than COSFORMER Other competing methods(Peng el al.) 2020 Choromanski et al., 2020; Kitaev et al.) 2019) are all based on kernel functions and have substantial performance gaps compared with our model. This validates the effectiveness of the proposed COSFORMER model compared with other efficient transformer variants.
3.4 RESULTS LONG-RANGE-ARENT BENCHMARK
To further evaluate the generalization ability of the proposed method, we train model from scratch on Long-range-arens benchmark 20206 Long-range-arena (Tay et al., (220b) is a bench- mark specifically designed for efficient transformers with long input sequences, thus serving a suitable testbed to assess the quality of efficient transformer variants comparatively. To ensure fair comparison, we first implement our method on Jax (Bradbury et al. 2018). then carefully follow their preprocessing, data split, model structure and training protocol. We evaluate our method on variety of tasks including Long sequence ListOps (Nangia & Bowman 2018). Byte-level text clas- sification (Maas et al. 2011), document retrieval using the ACL Anthology Network (Radev et al.) 2013), image classification on sequence of pixels on CIFAR-10 (Kizhevsky & Hinton 2009). and Pathfinder (Linsley al. 2018). As shown in Table 4 COSFORMER overall achieves competitive results across all the tasks while achieving best performance on ListOps and Document Retrieval. For the Pathfinder task, since the distance between the two points can be very far from each other, our introduced locality bias would have negative impact to this task and show a bit lags to other SOTA methods, despite that the performance gap between our method and the vanilla transformer is small It is worth mentioning that COSFORMER achieves the best overall scores on Long-range-aree benchmark, being one of the only two models that surpass vanilla transformer architecture.
7

---

Published as a conference paper at ICLR 2022
Table 4: Results Long-range-arena benchmark. Best result is in boldface and second best underlined. COSFORMER achieves the best average score across different tasks. Model
ListOps
Text?
Retrievalf
Imaget
Pathfinder
Avg 1
Local Attention (Tay et al. 2020b
15.82
$2.98
$3.39
41,46
66.63
46.
Lincar Trans. (Katharopoulor et
2020
16.13
65.9
53.09
42.34
75.3
50.55
Reformer KKitayy et al. 2019
37.27
56.1
53.4
38.07
68.5
50.67
Spanse Trans.(Child hild ct ai 2019
17.07
63.58
59.59
44.24
71.71
51.24
Sinkhorn Trans. Tay et 2020k
33.67
61.2
53.83
41.23
67.45
51.29
Linformen Wang
2020
35.7
53.94
52.27
38.56
76.34
51.36
Cerformer(Chorommann
2020
18.01
65.4
53.82
42.77
77.05
51.41
Synthesizer Tay et al. 2021
36.99
61.68
54.67
41.61
69.45
52.88
Longformer Bellagy ct 2020
35.63
56.89
42.22
53,46
Transformer Vaswani
2017
36.37
64.27
57.46
42.44
71.4
54.39
BigBird (Zaheer el al. 2020
36.05
$9.29
40.83
74.87
$5.01
COSFORMER
37.9
63.41
61.36
43.17
70.33
55.23
Table 5: Speed comparison on the long-range-arent benchmark in both training and inference varying
quence lengths (1-4k). We mark with cross method runs out of memory. The higher, the better.
Interence Speed per second)
Train (teedssteps per second)
model
IK
2K
3K
4k
IK
2K
3K
4K
Transformer Vaswani et al.
2017
25.37
7.83
X
X
6.95
2.23
x
x
Local Attention lay et al.
20200
57.73
33.19
23.36
17.79
13.45
6.71
4.32
3.09
Linformer Wang ct al. 2020
70.09
39.1
27.05
20.62
14.75
7.09
4.52
3.21
Reformer Kitaev el al. 2019
44.21
21.58
12.74
8.37
11.58
4.98
2.94
1.95
Sinkhom Trans. Tay cl
2020a
43.29
23.58
16.53
12.7
11.09
5.57
3.68
2.68
Synthesizer TTay
202
20.89
6.24
X
6.36
2.01
BirBird (Zaheor 2020
20.96
11.5
8.12
6.15
6.46
2.13
Linear Trans, (Kalharopoulos et al. 2020
67.85
38.24
26.28
11.86
5.54
3,53
Performer (Choromanski
74.15
42.31
22.44
6.49
2.94
Longformer [Bellagy 2020
22.99
6.72
X
4.4
Spurse Trans. Child 2019
24.87
X
6.77
COSFORMER
58.82
33.45
22.77
17.42
12.27
5.72
3.62
2.64
3.5 EFFICIENCY COMPARISON
In this section, we compare the efficiency of COSFORMER with other models, with a focus on long sequences as inputs. With the proposed linear attention module, we expect that COSFORMER scales comparably with other linear variants while significantly surpassing the vanilla transformer architecture. For a fair and comprehensive comparison, methods on Jax (Bradbury
implement method and competing al. 2018). We use the byte-level text classification benchmark and report runtime speed during both training and inference under different sequence lengths (Ik-4k). We conduct experiments on one Nvidia A6000 GPU and also report the corresponding inference- time memory foot prints shown in Figure T As shown in Table 5 and Figure [ most pattern based methods (Beltagy et al.) 2020 Zaheer et al.) 2020 Tay et af.) 2020a 2021) and vanilla transformer Vaswani al. 2017) are much slower and require greater memory than COSFORMER prevents them from extending to longer sequence. Further, the kernel based methods like (Narang et al.) 2021 Choromanski et 2020 Tay et al. (2020)) have comparable speed and memory overheads, but their performances are less satisfactory compared to COSFORMER across above metrics. In summary, our model COsFoRMeR achieves overall better efficiency than other linear variants while maintain superior modeling and generalization ability.
3.6 ABLATION: COS-BASED RE- GE-WEIGHTING
By introducing cos-based re-weighting. we provide non-linear mechanism concentrate the distribution of attention connections and stabilizes the training. this way, we encourage the model to better take into account the locality inductive bi- ases commonly observed many natural language tasks. In particular, we investi- gate the effect of the cos-based re-weighting in aspects. First, shown in Figure 4 by adding
Table 6: Performance comparison COSFORMER with and without cos-based re-weighting (ÓneLu). We evaluate on two compositive metrics. Bidirectional finetuneag: average score across 5 datasets reported in Table IRAngg average score across tasks reported in Table
Model
Bidirectional finetuneasy LRAws
dRaLU COSFORMER
85.12
54,20
85.25
55.23
8

---

Published as a conference paper at ICLR 2022
cos-based weighting, we obtain both notably better converge speed and final results in autoregres- sive language modeling. Further, in T we present a comparison between COSFORMER models with and without re- re-weighting mechanism. We use two composite metrics which comprchensively include 10 different datasets from bidirectional downstream fine-tuning tasks and long-range-arent Tay al., 2020b). COSFORMER achieves overall better results over the counterpart without weighting, improving the average scores on bidirectional finetuning and long-range-arena by clear margin. This verifies that the proposed re-weighting effectively incorporates the locality inductive biases for natural language tasks.
4 RELATED WORK
This section will introduce the existing works on improving the efficiency of Transformers, they can be broadly divided into two categories, Pattern based methods and Kernel based methods.
Pattern based method Pattern based methods sparsify the attention matrix with handcrafted or learnable patterns. As an early approach, Lee et al. (2019) leverages the inducing points from the sparse Gaussian process to reduce the quadratic complexities of a transformer. [Child et al. (219)) reduces the complexity by applying combination of strided pattern and local pattern to the vanilla attention matrix. Longformer (Beltagy et al., 2020 designs fixed diagonal sliding windows combined with global window, and the sliding window pattern can also be extended with dilation to enlarge the receptive field. Zahher et al. (2020) presents a more powerful and expressive sparse attention mechanism, which combines multiple types of attention patterns and gives a thorough study of sparse attention mechanism. Instead of fixed patterns, Kitaev et al. (2019) and Daras et al. (2020) group the attention computation process into buckets by local sensitive hashing, while |ooy et al. (2020) uses mini-batch spherical k-means. Nevertheless, Pattern based methods can only cope with sequences up to a certain length, and the computational complexity still grows rapidly when the input sequence becomes longer.
Kernel based method When faced with longer input sequences, it is more efficient to directly reduce the complexity of the theoretical calculation method. Kernel based methods speed up self- attention by reducing the computation complexity of self-attention from quadratic to linear. Vyas et al. (2020) approximate the full attention with a fixed number of cluster attention groups by assum- ing neighbouring queries in Euclidean space should have similar attention distributions. |eng al. (2020) chooses to use the production of Gaussian kernel functions to approximate Sofimax, chang- ing the order of scale dot product calculation, thus reducing the theoretical time to linear complexity and Choromanski et al. (2020) uses Haar measurement based kernel instead. Wang et al. (2020) imports the low-rank prior for attention matrix and approximate softmax with SVD decomposition manner. Xiong et al. (2021) utilizes the Nyström method with segment-means to generate a low- rank approximation of the Softmax matrix. Katharopoulos et al. (2020) formalizes the transformer layer as a recurrent neural network. In this paper, we demonstrate that the approximation to Softmux is unneccessary for Lincarization of self-attention module. We instead propose a new method to replace Softmax with linear operation with re-weighting mechanism, which reduces both time complexity and space complexity to O(N) while maintaining the accuracy.
CONCLUSION
We presented COSFORMER, a new efficient transformer that has linear time and space complexity. Our COSFORMER is based on two key properties of the original softmax attention: (i) every element in the attention matrix are non-negative, such that negatively-correlaeed information are not included for contextual information aggregation: (ii) the non-linear re-weighting scheme concentrates the dis- tribution of the attention matrix, in order to better exploit the locality inductive biases on sequence modeling. To fulfill these properties our COSFORMER, we utilized the RuLU function as our linear operation to ensure the non- non-negative property; new cos-based re-weighting mechanism was proposed to enforce the locality bias the original softmax attention. Since our COSFORMER nat- urally decomposable, it does not suffer the accumulated approximation error that usually happens in previous linear transformers. On causal pre-training, bidirectional pre-training, and multiple down- stream text understanding tasks, COSFORMER achieves comparable or even better performances than the vanilla transformer. On long sequence benchmark, COSFORMER achieved state-of-the-ee art performance over five different tasks. Further, COSFORMER obtains significant overall advan- tage in terms of time and memory efficiency over all existing efficient transformers, facilitating the transformers to easily scale to longer input sequence.

---

Published as a conference paper at ICLR 2022
REFERENCES
Abien Fred Agarap. Deep learning using rectified
linear
units
(relu).
arXiv preprint
1oXiv:/803.0885 2018.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling, In International Conference on Learning Representations, 2018.
Alexei Baevski, Yuhao Zhou, AbdeIrahman Mohamed, and Michael Auli, wav2vec 2.0: A frame- work for self-supervised learning of speech representations. Advances in Neural Information Processing Systems, 33, 2020.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05555 2020.
James Bradbury. Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman-Miln Jax: composable transformations of python+ numpy programs. Version 0.1, 55, 2018.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. arXiv, 2020.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European Conference Computer Vision, pp. 219-229. Springer, 2020.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint MxXiv:/904. /0509, 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 20K.::20009.14944 2020.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D Manning. What does bert look at? an analysis of bert's attention. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 276-286, 2019.
Djórk-Amé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). In Yoshua Bengio and Yann LeCun (eds.), 4th Inter- national Conference Learning Representations, ICLR, San Juan, Puerto Rico, 2016.
Giannis Daras, Nikita Kitaev, Augustus Odena, and Alexandros G Dimakis. Smyrf efficient atten- tion using asymmetric clustering. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), NeurIPS, volume 33. pp. 6476-6489. Curran Associates, Inc., 2020.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chupter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Pupers), pp. 1711-186, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An im- age is worth 16x16 words: Transformers for image recognition at scale. International Confer- ence on Learning Reprexentations, 2020.
Bolin Gao and Lacra Pavel. On the properties of the softmax function with application in game theory and reinforcement learning. arXiv preprint arXiv:1704 00805, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
10

---

Published as a conference paper at ICLR 2022
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizztion with gumbel-softmax. arXiv preprint arXiv: 1.001/11011100 144, 2016.
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers mns: Fast autoregressive transformers with linear attention. In International Conference on Ma- chine Learning, pp. 5156-5165. PMLR, 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2019.
Olga Kovaleva, Alexey Romanov, Anna Rogers, and Anna Rumshisky. Revealing the dark secrets of bert. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP- IJCNLP), pp.
2019.
A. Krizhevsky and G. Hinton. Learning multiple layers of features from tiny images. Master's 's thesis, Depariment of Computer Science, University of Toronto, 2009.
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation-invant neural networks. In JCML, pp. 3744-3753. 2019.
Drew Linsley, Junkyung Kim, Vijay Veerabadran, Charlie Windolf, and Thomas Seme. Learning long-range spatial dependencies with horizontal gated recurrent units. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pp. 152-164, 2018.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: /907. 1692, 2019.
Ze Liu, Yutong Lin. Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint arXiv:2103.1 /4030, 2021.
Antoine Liutkus, Ondfej Cifka, Shih-Lun Wu, Umut Şimşekli, Yi-Hsuan Yang, and Gaël Richard. Relative positional encoding for Transformers with linear complexity. In Marina Meila and Tong Zhang (eds.). Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 7067-7079. PMLR, 18-24 Jul 2021.
Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Y Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguisties: Human language lechnologies, pp. 142-150, 201
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. 5th International Conference on Learning Representations, ICLR, Toulon, France, 2017.
Nikita Nangia and Samuel Bowman. Listops: A diagnostic dataset for latent tree learning. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Com- putational Linguistics: Student Research Workshop, pp. 92-99, 2018.
Sharan Narang, Hyung Won Chung, Yi Tay, William Fedus, Thibault Fevry, Michael Matena, Kar- ishma Malkan, Noah Fiedel, Noam Shazeer, Zhenzhong Lan, et al. Do transformer modifications transfer across implementations and applications? arXiv preprint 11Xiv:2102.11992, 2021.
Jianmo Ni, Jiacheng Li, and Julian McAuley. Justifying recommendations using distantly-labeled reviews and Eine-graine aspects. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference Natural Language Processing (EMNLP-IJCNIP), pp. 188-197, 2019.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. In International Conference on Learning Representations, 2020.
11

---

Published as a conference paper at ICLR 2022
Dragomir R Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. The acl anthology network corpus. Language Resources and Evaluation, 47(4):919-94) 94-444 2013.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In J. Platt, D. Koller, Y. Singer, and S. Roweis (eds.), Advances in Neural Information Processing Systems, volume 20. Curran Associates, Inc., 2008.
Aurko Roy, Mohammad Taghi Saffar, David Grangier, and Ashish Vaswani. Efficient content-based sparse attention with routing transformers. In TACL, 2020.
Steffen Schneider, Alexci Baevski, Ronan Collobert, and Michael Auli. wav2vee: Unsupervised pre-training for speech recognition. In INTERSPEECH, 2019.
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. In arXiv, 2021.
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. In International Conference on Machine Learning, pp. 9438-9447. PMLR, 2020a.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. In International Conference on Learning Representations, 2020b.
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng, Synthesizer: Re- thinking self-attention for transformer models. In International Conference Machine Learning, pp. 10183-10192. PMLR, 2021.
Michalis K Titsias. One-vs-each approximation to softmax for scalable estimation of probabilities. arXiv preprint arXiv: 1609.07400 2016.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: An unified understanding for transformer's attention via the lens of kernel. In EMNLP, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6000 2017.
A. Vyas, A. Katharopoulos, and F. Fleuret. Fast transformers with clustered attention. In NeurIPS, 2020.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Glue: A multi-task task benchmark and analysis platform for natural language understanding. In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 353-355, 2018
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint 0407022006604768, 2020.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Víkas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. In AAAI, 2021.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences, In NewrPSS, 2020.
12

---

Published as a conference paper at ICLR 2022
A APPENDIX
A.1 MATHEMATICAL DERIVATION OF cOS-BASED RE- NE-WEIGHTING
Following Equation
we give a detailed deviation of how to obtain output at position ith position:
0, Σn F(Q,,
E==1 s(Q,,K)
+ Quinn (Rgen) I
1=
@ia0 (Rgae)²
Vg
ΣØ=1
N
@:000 (Rgoa) T
+ Quin (Kgan)
T
(13)
2==1
N
2:o0 (Ko00) T
Vg + 2j=1 N
Quinn (Kgin)
V3
Σj 2:00 (Kooa)² T
+ + Quinn (Rgin)
j= Q: (Kgan))
T V3) + Σラ= Quin
((Rgtt) T
Vi
2jj11 Q000 (Kgun)²
+ Σソ N j=
Qrin
Kgin
K
where i,j 1, N, M 2 N, and Q' (c) K' RcLU(K). Let Q = Q cos (2M ) xi
Q = Q cos 2M K = K2 cos 2M ). K K2 sin ( ). It presents that the output of the proposed CORFORMER attention be obtained in linear manner.
A.2 PSEUDO CODE OF COSFORMER
Algorithm describe the way to compute COSFORMER attention
Algorithm 1 COSFORMER attention
Input: Q E Nxd1 K E R M V E R Mxds;
Output: O E Nxds;
Use M, to represent the i-th row of matrix M;
Initialize Ali
O
0, 1,....
N.j 1,...,dy; d2:
2N
Initialize 5000
[
=
0, S S
ÍÍ = 0, T
= 0, in
[1] = 0, 1,...,4,,3 = 1....,dy
for i in 1,..., M do:
Koos K; cos
Ti
Kain = K2 sin
2M ):
JM
S += (Kg00)
V::
Ssin +- Krin)
T
V::
Teos Koos.
Trin Kiim;
end for
for i in 1, N do:
Q006 = Q: i cos xi Qthn = Q sin 2M
2M
0 Q goos goin
end for
A.3 ALGORITHM TO VISUALIZE ATTENTION MATRIX
Algorithm 2 describe the way visualize attention matrix Figure
13

---

Published as a conference paper at ICLR 2022
Algorithm 2 Algorithm to visualize attention matrix
Input: M, E BRdd k 1,,,.,
ni; threshold E [o, 1];
Output: M Rdxd.
Initialize M 0,i I.....,,,d77 ..
for in :......1. do:
for i in 1...... d do:
index = argsont Mxli)) (in descending order)
p = 0 for j in I..... do:
index []
p +- M4 [1]
[+][[]] +=1
if p >> Uhreshold then: break
end if
end for
end for
end for
M /- n
Use heatmap to visualize M;
A.4 INTRODUCTION OF DATASET
We train both models on autoregressive language modeling and bidirectional modeling by Wikitext- 103 dataset, it split by tokens and its statistics as Table Then we fine-tune the pre-trained bidi- rectional modeling on several text classification tasks.
QQP dataset contain thousands of sentence pair from community question-answeringg website Quora.Network need to determine pairs of question are semantically equivalent. SST-2 and IMDB are collections of movie reviews. The task is to determine whether a review is positive or not. AMA- ZON dataset contains millions of product reviews from Amazon. The requirement of this task is to infer the scoring of the product from the review text.MNII is a crow-source collections of sentence pairs. The network must distinguish which of the three categories entailment, contradiction and neutral the given sentences belong to.
The long-range-aren benchmark contains different datasets. List contains some designed clever mathematical problem to clarify the parsing ability of neural models. IMDB is also used in this benchmark to examine the text classification ability of neural models. CIFAR-10 is image col- lection of various of object, this task require models capture 2D spatial relations between flatten pixels.In In pathfinder task, models need to determine the connection two points in the picture, so as to examine the model's ability to acquire 2D spatial relationships.AA. dataset is used to evaluate the ability for models to encode and store compressed representations for retrieving.
Data
Train
Valid
Test
WikiText-
103
103M
218K
246K
QQP
364K
391K
SST-2
67K
1.8K
MNLI
393K
20K
IMDB
25K
25K
AMAZON
3M
168K
168K
ListOps
90K
10K
AAN
147K
18K
17K
CIFAR-10
50K
10K
Pathfinder
160K
20K
Table 7: Statistics for the datasets.A subset of "Small" amazon subset on electronics category is used for experiment
14

---

Published as a conference paper at ICLR 2022
A.5 QUALITATIVE RESULTS OF LRA
We provide our qualitative results of the ListOps and Document Retrieval tasks on Long-Range Arena benchmark (Tay et al. 2020b) with comparison the vanilla transformer.
ListOps is a ten-way classification task which aims to prediction the results sequence with a hierarchical structure and operators MAX, MEAN, MEDIAN and SUM MOD that are enclosed by delimiters (brackets). The network needs to access all tokens and model the logical structure of the inputs in order make prediction.
Document Retrieval task is to decide whether the two input long documents are similar or not with a binary label. This task evaluates model's ability encode and store compressed representations that are useful for matching and retrieval. Since the samples in LRA are too long, We substantially shorten some selected samples and display them as below:
Listops:
Input: 8
[MND ]
9
3
6 Our Output: 0,
2
) 8
2
Transformer output: 9,
0
3 Input: ( (((((11
ISM 5
1
( (. (...... (
Input: ( ( [MIN 5 1 0 ( IMED 8 7 2 8
8 1 7 11l) Our output: 9, Transformer output: 3,
Ground-truth: 9
5 Input: ( (C((0(6 [MAX 7 8
[MAX 5
) ( ( ( [SM 3 6 (
1 6 4 2
110D Our output: 9, Transformer output: 5, Ground-truth: 9
Listing 1: Examples of LisOps
Byte-level document retrieval:
I Textl: b'1 Introduction Recent advances statistical Machine Translation (SMT) are widely centred around two concepts: (a) hierarchical translation processes, frequently empleying Synchronous Context Free Grammars (SCFGs) and (b) transduction synchronous rewrite processes
linguistic ......
3 Text2:
b'1 Introduction Automatic Grammatical Error Correction (GEC) for non-native English language Learners has attracted more and more attention with the development of natural language processing machine learning and big-data techniques. ?The CoNL12013 shared task focuses on the problem of GEC in five different error types including determiner, preposition, noun number....
4
S Cur output: False, Transformer output: True, Grourd-truth: False
Listing 2: Examples of Document Retrieval
15


Reviews:
[{"review_id": "zMuI54qpv5d", "reviewer": "Reviewer_yo75", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "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+ The paper is very well written and organized. It explains the problem and the proposed solution clearly.\n+ The authors have been thorough in experimental analysis, and have shown the results on benchmarks as well as the generalization ability of the method. \n+ The idea of reweighing with a cos based function is faily novel. \n\nWeaknesses\n- In Section 3.2,  it is mentioned that \"in typical natural language tasks, the feature dimension of one head is always much smaller than the \n input sequence length N\". What implications it could have when d is not that small. Would the analysis will still be applicable (especially the computational speedup for domains such as images) ?\n-  In Table 4, results on \"Pathfinder\", the method significantly lags behind the state of the art (~10%). An explaination would help understand the challenges of this particular setting. \n- As the method is shown to significantly outperform state of the art on a few benchmarks, a few examples on where it performs better while the other methods perform inferior should be provided to qualitatively analyse the improvement. \n- A few typos (i) Section 3.1 \"...the Eq. 2 become\" -> Eq. 2 becomes (ii) Section 3.4, \"...such locality bias, ie.,...\" -> i.e. (note the dots)\n", "summary_of_the_paper": "The paper proposes a variant of the transformer network. The authors base their work on their analysis of softmax attention. They argue that softmax works well because of two reasons (i) it stabilizes training due to reweighing the attention as well as associated connections within the network, and (ii) it forces non negative values in the attention matrix. Utilizing these observations, they propose two modifications (i) a linear projection kernel i.e. ReLU to compute similarity, and (ii) a cosine function for reweighing the attention values. With exhaustive experiments on 5 benchmarks, authors demonstrate that the proposed modifications lead to state of the art performance with reduced computational compelxity.   ", "summary_of_the_review": "The paper is generally well written and the problem is motivated clearly. The paper also highlights contribution with reference to the compared methods. The experiments (empirical) are exhaustive, however, a few qualitative results on explanations of a few disparities in the performance of the proposed method (as indicated in the detailed review) would help understand the impact of the proposed method better. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "t2QMCxSaAyN", "reviewer": "Reviewer_ZpU1", "summary": "", "questions": "", "limitations": "", "rating": 6, "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\n1. The proposed method accelerates the inference speed of the model.\n\n## Weakness\n\n1. The motivation of using ReLU is unclear. Though softmax promises a non-negative weight, there is no evidence showing that non-negative weights is essential.\n\n2. The motivation of cosine operation is unclear. The success of self-attention partly attributes to its ability of building long-distance dependency. The authors introduce locality into self-attention through consine operation but do not show its rationality.\n\n3. Lack of essential ablation studies. How about the model without ReLU?\n\n4. Limited improvement. In terms of Table 6, the performance gain induced by the cosine operation is limited (~0.13) on some models.", "summary_of_the_paper": "The paper proposes a substitute of the vanilla self-attention module, which claims a linear complexity. Like Performer and many other previous works, it introduces kernel functions and change the computation order among QKV in the self-attention module. Essentially, it replaces softx with ReLU and a cosine operation. Experiments are done on several benchmarks.", "summary_of_the_review": "Totally, I think the paper is still not well-prepared for publication.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "nW6oZqFxHUS", "reviewer": "Reviewer_iBUe", "summary": "", "questions": "", "limitations": "", "rating": 8, "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- The proposed method is simple and intuitive and experimentation on both LRA and real world language modeling task shows that it provides better speed-accuracy trade-offs when compared against other effective attention mechanisms.\n- The paper is well written and easy to follow.\n\nSome questions and concerns:\n\n- Related work section could be expanded to include discussion on other efficient attentions such as:\n    - Approximating softmax atttention:\n        - SMYRF: Efficient Attention using Asymmetric Clustering\n        - Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention\n        - Fast Transformers with Clustered Attention\n    - Sparsity pattern\n        - Efficient Content-Based Sparse Attention with Routing Transformers\n    - Low rank projections:\n        - Set Transformer: A Framework for Attention-based Permutation-Invariant Neural Networks\n- Related work on the properties for kernelized attention (non-negative similarity scores):\n    - Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel\n\n- Clarification:\n    - For bidirectional modeling, why are the models pre-trained only to 50K iterations? From figure 4, the validation perplexity still seems to be decreasing especially for the vanialla attentions.\n    - In the table 1, is $\\phi_{ReLU}$ linear attention with ReLU as the non-linear activation to get positive similary scores?\n    - Similarly for table 6, without re-weighing is linear attention with ReLU activations?\n\n- Relation with relative positional encodings for linear attentions:\n    - The cosine re-weighing can also be looked as an efficient way to introduce relatve positional bias to the linear attention. In this regard it would be good to contrast this to recently introduced relative positional encoding for linear attention such as (a) Rotary Position Embedding and (b) Relative Positional Encoding for Transformers with Linear Complexity       \n\n- Page 5: \"The non-linear re-weighting mechanism introduced by the softmax attention can concentrate the distribution of the attention weights and therefore stabilize the training process. It also means that, to some extent, it will punish far-away connections and enforce locality.\" While this is empirically observed, softmax normalization doesn't introduce such inductive biases. It might be better to rephrase this differently.\n", "summary_of_the_paper": "This work introduces a simple and elegant way to incorporate recency bias for kernelized linear attention mechanism using cos-based re-weighting mechanism.  This cosine re-weighting can be effectively implemented via query and key matrix transformations for linear attention mechanisms. The authors experiment with LRA, and both language modeling with autoregressive and bidirectional setups to show the efficacy of proposed method.", "summary_of_the_review": "This work introduces a simple mechanism to add relative attention / locality bias to linear attention to significantly improve their performance. The results showcase that using the proposed cosine re-weighting the linear attention achieves similar performance to vanilla transformers while being significantly faster.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "bKXTS2xh-ZI", "reviewer": "Reviewer_M72S", "summary": "", "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": "**Strengths**\n\n- The paper is well-written. With a clear logic flow, the paper is also easy to understand.   \n\n- Two factors for improving the efficiency of transformers: (1) the attention matrix should have non-negative elements and (2) aggregating negatively-correlated information needs the non-linear re-weighting strategy, are reasonable, especially the latter one.\n\n- The proposed cosForm, leveraging ReLU and  the Ptolemy’s theorem (cosine re-weighting), is simple yet new, to the best of my knowledge. Glad to see it shows a good accuracy-efficiency tradeoff when applying it to replace the softmax operator in calculating self-attention for transformers. \n\n**Weaknesses**\n\n- In the Method section, when analyzing the computational complexity for the vanilla self-attention and the linearized design, the authors assume \" The input length is N and feature dimension is d, with d << N.\" However, in the experiments, such an assumption does not well hold. For instance, in section 4.1, the sequence length is 512 while the projected feature dimension is 1024. In this perspective, I am concerned by the claimed motivation/advantage of the proposed method, can \"linear space and time complexity\" still be well supported? \n\n- The limitations of cosForm should be discussed, particularly a more in-depth discussion of its connections to many existing works is necessary. \n\n- As for the comparison of runtime speed in training and inference on the long-range-arena benchmark, did the authors adopt the same transformer architecture for all methods? Why different methods sometime show distinct trends for runtime speed in training vs. inference, e.g. Linear Trans. vs. cosFormer?\n\n- How about the performance of applying cosFormer to other tasks such as in computer vision besides NLP tasks? \n\n- Some grammar errors or typos:\n\n   on page 2, \"Kitaev et al. (2020) group\"->\"Kitaev et al. (2020) groups\"\n\n   on page 4, \"a computation complexity of O(Nd)\"->\"a computation complexity of O(Nd^2)\"\n\n   on page 7, \"despite Longformer (Beltagy et al., 2020) achieve\"->\"despite Longformer (Beltagy et al., 2020) achieves\"", "summary_of_the_paper": "This paper presents cosFormer which presents a linear operator to replace the softmax operator in calculating self-attention for transformer architectures, maintaining competitive predication accuracy while enjoying linear space and time complexity instead of quadratic costs compared to the vanilla baseline. Under this context, the authors emperically show that two factors are critical to improve the efficiency of transformers: (1) the self-attention matrix should have non-negative elements and (2) aggregating negatively-correlated information needs the non-linear re-weighting strategy. Based on the insightful analysis, the proposed cosFormer adopts ReLU and the Ptolemy’s theorem (cosine re-weighting) in designing a linear replacement of the softmax operator. The performance of cosFormer is validated with the language modeling (autogressive in both causal and non-causal cases) on WikiText-103 dataset,  the finetuning on a lot of downstream datasets such as GLUE, IMDB and AMAZON, and the comparison with state of the art methods on the long-range-arena benchmark.  ", "summary_of_the_review": "Please see the comments in \"Main Review\".", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces a new linear attention mechanism for transformer based models.  This is accomplished by replacing the softmax in the standard transformer self-attention with a cosine-based re-weighting mechanism.  The empirical results are good, and cosFormer generally outperforms existing efficient transformers for autoregressive language modeling, fine-tuning, and on the long range arena.

The reviewers were generally positive regarding the paper, with all reviewers voting to accept.  The discussion period focused on particular choices regarding the ReLU activation function vs. other non-negative activation functions, further motivating the cosine operation, and comparing the speed of cosFormer vs. other efficient transformers.  The authors responded by providing additional ablations to empirically validate the choice of ReLU, motivated the cosine operation by noting that it introduces a locality bias, and further described the computation requirements of their transformer vs. prior work.

Overall, this is an interesting addition to the linear / efficient transformer literature, with solid empirical results supporting the various design decisions.
Decision: Accept (Poster)