SEA: Sparse Linear Attention with Estimated Attention Mask
Authors: Heejun Lee
Jina Kim
Jeffrey Willette
Sung Ju Hwang
Venue: ICLR 2024 poster (2024)
Keywords: Sparse Attention
Attention Estimation
Linear Attention
Transformers
NLP
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: We present SEA, a linear attention method performing sparse attention in linear time with a generated attention mask based on compressed attention matrix estimation, using kernel-based linear attention.

Abstract: The transformer architecture has driven breakthroughs in recent years on tasks
which require modeling pairwise relationships between sequential elements, as
is the case in natural language understanding. However, long seqeuences pose a
problem due to the quadratic complexity of the attention operation. Previous re-
search has aimed to lower the complexity by sparsifying or linearly approximating
the attention matrix. Yet, these approaches cannot straightforwardly distill knowl-
edge from a teacher’s attention matrix, and often require complete retraining from
scratch. Furthermore, previous sparse and linear approaches lose interpretability
if they cannot produce full attention matrices. To address these challenges, we
propose SEA: Sparse linear attention with an Estimated Attention mask. SEA
estimates the attention matrix with linear complexity via kernel-based linear at-
tention, then subsequently creates a sparse attention matrix with a top-k̂ selection
to perform a sparse attention operation. For language modeling tasks (Wikitext2),
previous linear and sparse attention methods show roughly two-fold worse per-
plexity scores over the quadratic OPT-1.3B baseline, while SEA achieves better
perplexity than OPT-1.3B, using roughly half the memory of OPT-1.3B. More-
over, SEA maintains an interpretable attention matrix and can utilize knowledge
distillation to lower the complexity of existing pretrained transformers. We be-
lieve that our work will have a large practical impact, as it opens the possibility of
running large transformers on resource-limited devices with less memory.
Code: https://github.com/gmlwns2000/sea-attention

Published as a conference paper at ICLR 2024
SEA: SPARSE LINEAR ATTENTION WITH ESTIMATED ATTENTION MASK
Heejun Lee', Jina Kim Jeffrey Willette,, Sung Ju Hwang?
School of Computing', Graduate School of A1 Korea Advanced Institute of Science and Technology!:2 DeepAuto.ai
Daejeon, South Korea {ainl, jinakim, jwil swillette,
ABSTRACT
The transformer architecture has driven breakthroughs in recent years tasks which require modeling pairwise relationships between sequential elements, as is the case in natural language understanding. However, long seqeuences pose problem due to the quadratic complexity of the attention operation. Previous re- search has aimed to lower the complexity by sparsifying linearly approximating the attention matrix. Yet, these approaches cannot straightforwardly distill knowl- edge from teacher's attention matrix, and often require complete retraining from scratch. Furthermore, previous sparse and linear approaches lose interpretability if they cannot produce full attention matrices. To address these challenges, we propose SEA: Sparse linear attention with an Estimated Attention mask. SEA estimates the attention matrix with linear complexity via kernel-based linear at- tention, then subsequently creates a sparse attention matrix with a top-k selection perform sparse attention operation. For language modeling tasks (Wikitext2), previous linear and sparse attention methods show roughly two-fold worse per- plexity scores over the quadratic OPT-1.3B baseline, while SEA achieves better perplexity than OPT-1.3B, using roughly half the memory of OPT-1.3B. More- over, SEA maintains interpretable attention matrix and utilize knowledge distillation to lower the complexity of existing pretrained transformers. We be- lieve that our work will have a large practical impact, as opens the possibility of running large transformers on resource- limited devices with less memory.
INTRODUCTION
The transformer (Vaswani et al., 2017) architecture has revolutionized various fields of artificial telligence, such natural language understanding (Touvron et al., 2023; Wang et al., 2022) and computer vision (Dosovitskiy et al., 2021) due to its ability to learn pairwise relationships between all T tokens in given sequence (O(T2).. This has ushered in the era of large transformer-basel foundation models with impressive generalization abilities (Brown et al., 2020; Chiang et al., 2023). However, since the transformer's attention mechanism comes with quadratic space and time com- plexity, it becomes untenable for handling long sequences which are essential for tasks such as di- alogue generation (Chen al., 2023). To overcome this limitation, previous works have suggested approaches with linear complexity by using static or dynamic sparse attention patterns (Beltagy al., 2020; Zaheer et al., 2020; Tay et al., 2020a; Kitaev et al., 2020: Tay et al., 2020b; Liu al., 2021), or by replacing quadratic attention with kernel or low-rank approximations (Choromanski al., 2021; Chen et al., 2021; Qin et al., 2022).
However, despite their promising aspects, previous linear attention methods have yet to be widely used in research and production for several reasons. Firstly, these attention mechanisms cannot be easily swapped into existing finetuned models. After radically replacing the quadratic attention mechanism, they require training new attention relations to attempt to recover the performance of the quadratic attention module and cannot directly leam the full range of attention relations during knowledge distillation. Therefore, they suffer from unpredictable accuracy degradation on down- stream tasks. For example, as shown in Table A.8, Reformer (Kitaev et al., 2020) outperforms many

---

smaller memory budget as shown in Published as a conference paper at ICLR 2024
Adatcion Reutionship Knowledge Transfer
Vanilla Attention (teacher) 0
Previous Works
Perfemerr
Insepdoo Spare Masc
kooodedge
Trrater
KT
<V Clyi
<(Q)xx
XXX
Consert Bky Spase Wesk
Sparse Attention O(2)
Kernel-based O(T)
SEA attention (Ours) O(T)
Figure I: Concept. We estimate the attention matrix in compressed size (A), then perform a grouped top- k selection, and subsequently perform sparse attention with novel FlatCSR operation using an estimated attention mask on the full attention matrix. SEA has linear complexity in all steps at test-time, and requires direct attention matrix distillation from the quadratic teacher at train-time.
other baselines MNLI (Williams al., 2018); however, it shows the worst performance in Ta- ble 2 on Wikitext2 (Merity et al., 2017). Secondly, previous linear attention methods may hinder the ability interpret the attention matrix or merge/prune tokens (Kim et al., 2022; Lee al., 2023; Bolya et al., 2023) of the transformer model, as they do not produce the full attention matrix which usually required for analyzing the importance of a given token.
In contrast to previous works, our proposed linear attention method, SEA: Sparse linear attention with Estimated Attention mask, focuses on estimating the attention matrix from a pretrained teacher transformer model with OTT) complexity at inference rather than ( In order to do so, our novel estimator, distills knowledge (Hinton et al., 2015) from the teacher and estimates the attention matrix with (TT) complexity by fixing the second dimension to constant value K where K < T. This results in compressed attention matrix which can be decompressed via interpolation into an approximation of the full attention matrix when performing the distillation step as opposed to previ- ous works which prescribe retraining attention relationships during distillation due to incompatible changes in the attention operation (Choromanski et al., 2021; Qin al., 2022). By distilling from the full attention matrix into our compressed matrix, SEA can use the complex and dynamic attention information from the pretrained model and preserve task performance.
Furthermore, as SEA distills knowledge from a full attention matrix, the resulting compressed at- tention matrix and sparse attention matrix can still provide interpretable attention by allowing for interpreting the relationships and importance of tokens (e.g. analysis of attention patterns in im- ages (Dosovitskiy et al., 2021) and token pruning (Kim et al., 2022; Lee et al., 2023; Bolya al., 2023)). Lastly, SEA reduces the space and time complexity of attention from O(T²) into O(T) test-time, with significantly reduced memory and computational cost while maintaining similar performance to the original pretrained transformer model, as shown in Fig. 5 and Tables 2 and A.8.
In Fig. 1, we provide a overview of our proposed method's inference stage. The main idea is to first estimate a compressed attention matrix Á with (TT) complexity, and subsequently perform sparse attention after choosing only (T)) important relationships inferred from Å. Specifically, we decode the output features of the kernel-based linear attention method Performer (Choromanski al., 2021) to form Å. Next, we perform a top-k selection on Å to form compressed attention mask which can be used to generate sparse attention mask for the final sparse attention operation. By utilizing both kernel-bast based and sparse attention, we can take advantage of their diverse token embedding spaces, as shown in previous work (Chen et al., 2021). The compressed estimated attention matrix Å is trained via knowledge distillation (Hinton et al., 2015; Jiao et al., 2020) from a pretrained quadratic teacher model to distill complex and dynamic attention patterns. We achieve linear complexity in this process by controlling sparsity and compression ratio in the compressed attention mask.
We validate SEA by applying to BERT for text classification (Devlin et al., 2019; Wang et al., 2019) and to OPT for causal language modeling (Zhang et al., 2022; Merity et al., 2017). Our em- pirical findings demonstrate that SEA adequately retains the performance of the quadratic teacher model in both tasks (Tables 2 and A.8), while previous methods do not. SEA significantly out- performs the best linear attention basclines, such as Performer, in language modeling with 47.7% perplexitywiiin
lower (better) perplexity while consuming 54.2% less memory than quadratic attention (Table 2). This opens up the possibility of running long context language models on lower- lower-grade computing devices that have smaller VRAMs because SEA can Section 5 and Fig. 8. To summarize:
2

---

Published as a conference paper at ICLR 2024
We propose a novel, test-time linear attention mechanism (SEA) that distills knowledge from a pretrained quadratic transformer into compressed estimated is then used to create a sparse attention mask for the final attention operation. As demon- strated in Fig. 8, SEA is (T) at test time, where there is no distillation step.
estmmaatd attention matrix which
We demonstrate the efficiency and effectiveness of our method through empirical eval- uations on natural language processing tasks such as text classification on GLUE and language modeling Wikitext-2, where we maintain competitive performance with the vanilla transformer baseline, as shown in Figs. 5a and 7a.
We propose and provide code for a novel CSR tensor operation, called FlatCSR, which is capable of handling non-contiguous flatten tasks within a GPU kernel.
We showcase the interpretability of our method by visualizing the estimated sparse atten- tion matrix and comparing to the teacher's attention matrix. Our estimator can estimate both self-attentio and causal attention.
2 RELATED WORK
Sparse Attention for Efficient Transformers. Sparse attention can be categorized into methods using static and dynamic attention masks. For static attention masks, Longformer (Beltagy et al., 2020) and BigBird (Zaheer et al., 2020) introduce heuristic patterns, but since token relationships may not fit to those heuristics, it can be challenging to achieve state-of-the-att performance for every task. Hence, some recent methods focus on learning to sort or cluster the tokens to better fit static masks (Tay et al., 2020a; Kitaev et al., 2020; Tay et al., 2020b). However, these works still operate based on static patterns, a more flexible and learnable setting is necessary since patterns in attention unavoidably dynamic and data-dependent, as shown in TDSA (Liu al., 2021), which learns to estimate an attention mask. However, TDSA still performs quadratic attention when generating dynamic sparse attention masks.
Kernel and Low-rank Methods for Efficient Transformers. Recent works either focus on using bernel-based methods (Choromanski et al., 2021; Qin et al., 2022) or combining kernel and sparse methods (Chen et al., 2021). For example, Performer (Choromanski et al., 2021) uses positive- definite kernel approximation to approximate the softmax attention mechanism but comes with non- negligible approximation errors and therefore is not generalizable on every task. To handle this problem, Cosformer (Qin et al., 2022) replaces the non-linear approximate softmax operation of Performer with linear operation that contains non-linear cosine-based re-weighting mechanism. But again, the cosine weighting scheme is a heuristic which limits it to approximating attention with a specific pattern, which may be a downside for attention matrices which do not follow the fixed heuristic pattern, as shown in Fig. 9 (subfigure a, bottom-left). Another work, Scatterbrain (Chen al., 2021) proposes combining sparse and low-rank methods to approximate attention, by separately applying both and then summing the result together. While this is a promising approach, it still cannot straightforwardly benefit from direct attention matrix distillation.
SEA: SPARSE LINEAR ATTENTION WITH ESTIMATED A ATTENTION MASK
Preliminaries. We first define notations used for the attention mechanisms. We define the following real matrices: Q K, V E Txx, The attention matrix A is defined as A softmax softmax(QQ"") e be the context feature of the attention lookup, and T be the length of the input sequence, and d be the size of the head embedding dimension. We elementwise multiplication which may be applied between similar sizes matrices or on the last dimension and repeated over prior dimensions as is the case when applied to a matrix and vector. We define K as the width of the compressed matrix, k as a value which controls the sparsity of an attention matrix. The matrix superscript
RTXT Let C = AV e Txx
@ denote
(e.g. A*) indicates a sparse matrix, and superscript" (e.g. A) refers to the compressed T x K space. Let (L, q) and MSE(z, y) be the standard KL divergence and mean-squared error, respectively. All matrices may also contain a batch (B) and head (If) dimension, but we omit this for simplicity unless otherwise specified.
Performer (FAVOR+). We define the simplified notation of FAVOR+ (Choromanski et al., 2021). We omit details, such as the kernel projection matrix, for the convenience of reading. The output of FAVOR+ is defined FAVOR+(Q, K, V) o(Q) x (ø(K) (TVo(V)) R As there is no need to construct the full T x T attention matrix, the FAVOR+ mechanism scales linearly.
3.1 SEA ATTENTION
SEA consists of two main phases that are depicted in Fig. 2: (1) Attention Estimation (kemnl- based), and (2) Sparse Attention Mask Generation and Subsequent Sparse Attention. In step
3

---

Published as a conference paper at ICLR 2024
Teacher
Distill Teacher Knowledge to SEA
Deploy SEA
Transformer Block
Transformer Block
(1) Attention Estimation
MLP
Decoder
MLP
Performer (incade Contert
MUP
(revaive))
MHA
CNN
X*
Sparse Actention Mask Generation
Sparse Attention
SEA Attendion)
ynterolais Calmmaionn for
A Presched
Teacher Attention Knowledge Distillation (KD)
Figure 2: Model diagram of our proposed linear attention method, SEA. SEA replaces the multi- attention (MHA) mechanism from the teacher transformer model while preserving the attention matrix using knowledge distillation with only a small finetuning dataset. The deployed SEA model shows linear complexity at test time.
(I), we estimate a compressed attention matrix Á by fixing the size of one dimension to be K, where K < 7 using the kernel-based method Performer and a simple decoder. In step (2), we build an attention mask from the values in Á using one of our novel grouped top-k selection methods depicted in Fig. 4. We then interpolate the mask to form sparse mask on the full attention matrix A. resulting in sparse attention matrix A". Then, we perform the sparse A'V multiplication to complete the procedure. By using kernel-based estimator, well as sparse attention matrix, we take advantage of the complementary aspects of their representation spaces, as shown in previous work (Chen al., 2021). Detailed model structure is shown in Fig. A.6.
Attention Matrix Estimation. The intuition for Å is to compress one dimension of the matrix in a way analogous to compressing one dimension of an image. We ultimately want the values within the compressed attention matrix be similar to the full attention matrix in the same way that compressing one dimension of an image creates a distorted image which is semantically similar to the full image. For this task, we use the Performer (Choromanski et al., 2021) and a decoder, and treat the output as our estimated attention matrix. Once this attention matrix 'image" is obtained, it can be interpolated to decompress into an approximation of the full attention matrix for distillation during training, or likewise, we may also perform top-k selection on the compressed matrix in order to construct a mask for sparse attention. Our attention matrix estimation starts by passing Q, K, Vea to the kernel-based linear attention method, Performer (Choromanski et al., 2021), where Van = =Vv;; V) E R and V, is obtained by performing nearest neighbor interpolation on the identity matrix I E d to modify the first dimension resulting in matrix Vi E R 7xxd We include V, in Vest because the output of Performer can be considered the attention matrix when V =IE
RT'X" (e.g. Q(K I) QK ), as shown by Choromanski et al. (2021). Therefore, passing VI together with V may enable a more accurate estimation of the attention matrix by the decoder
described in the next section. This results in the context encoding C Pert FAVOR+(Q2 Q, K, Vea) E RT22
CNN Decoder. We further transform Performer's estimated out- put Cpen with an MLP and CNN. We found the CNN to be nec- essary part of SEA due to the fact that convolutions provide fine- grained detail among local features, which is crucial for dynam- ically representing complex patterns present in the attention ma- trix shown in Figs. 3 and 9. As we may consider the atten- tion matrix as a kind of compressed 'image,' a CNN is a natural choice for the decoder. For the decoder, we begin by concatenating the local Performer encoding Openf and the previous context V as
Octput
CNN
Figure 3: Visualization of Input and Output of CNN Decoder
V (Cpoti V] E Rxx3d We then apply an MLP μ 3d + Rd to obtain intermediate representation Z = (((', Txxx' where d' is a hyperparameter which decides the shared hid- den state size. The resulting Z is used for estimating the attention matrix and also for the scalers (Smix and Sproh) described in Section 1.1.1. Then, before applying the CNN on Z, we apply MLP v: I2d + +*RKas Ken///, where and C₁ are respective width reduction and channel expansion factors. We transpose and reshape to make the output 2 = v(Z) into He,, x xTx k/c,, Finally, we apply a 2D CNN fae: which treats the extra head dimension H a channel dimension, resulting in the compressed attention matrix Á face (2) E T (for further details, see Appendix A.5.1). As the decoder fae: results in a fixed width Á, we can successfully generate dynamic patterns with linear time and space complexity. The CNN faee plays significant role due its ability to capture
4

---

Published as a conference paper at ICLR 2024
local patterns of the estimated compressed attention matrix. The depth of the CNN can be adjusted depending on the task, but we use a fixed 3-layer CNN in our experiments.
Grouped Top- Selection. Once we have the compressed attention matrix A, we must select crit- ical values which will be used in the final T x 7 sparse attention. However, since our top-k selection is held in the compressed A, we transform A to k, which refers to the number of selected values in the compressed T x K attention space. For this, we propose four novel methods for top-k selec- tion: per-query, per-head, per-batch, and causal-per-batch where each method performs the top-k selection along different dimensions, as shown in Fig. 4. As noted in Section 3.1 (preliminaries), we have previously omitted the head dimension H from all matrices. However in this paragraph, for clarity, we include the head dimension H so that Å E R per-batch, we gradually extend the dimension of the group, starting from the last dimension of A, so that top-i selection held in RK, IR 1R7*K K and IR xxTx K space for each grouping method, respec- x h. and kpectuch Finally, we propose causal- causal--er-batee per-batch for causal attention, which performs top-k in # space, with Recasol--er.hatch H x k. For this, we transpose Á to Tx and group the last two dimensions without the T dimension, to avoid temporal information exchange across the time dimension. In our experiments, we use -crsaal-er-batch which shows strong performance in our ablation study on GLUE-MNLI (K- = 128, 5 epochs), as shown in Table 1.
For per-query, per-head, and
tively. Consequently, & is also adjusted to hiper-uury = = hipereead
x T x k.
Linear Sparse Attention Mask Generation. With the obtained
Table 1:
Ablation study on grouped top-k modes
Å, we generate a sparse formatted binary mask M E {0, 1)7XT, For this, we take the following two steps: 1) Performing our pro- posed grouped top-k selection from the compressed Á to gener- ate a binary mask M E ( and 2) interpolating M to the sparse formatted M E {0, 10,1"" For the grouped top- k selection, we set k as hyperparameter, which will determine the number of selected indices in each block of the binary mask M' depending on the top-k strategy and the attention matrix size. Note that each selection strategy has a different block (group) shape as depicted in Fig. 4. However, since we perform top-k on the smaller Á E TTTT,, we must convert & to a compressed & with k = max(1, 11, round * K/T)). Once we obtain the compressed mask M f0. 1)7xK we interpolate it into the sparse formatted x2" that max(1, rounddk= terpolation (pixel duplications) to create M will become a function of 'T' and no longer have linear complexity, as this results in a block larger than being computed for M Therefore, in order to avoid this, we enforce the number of pixel duplications in the block of M to be min(k, [T] K1), and uniformly space the resulting pixels within the larger block. Since we only need to check the indices where the values are in the compressed M and put I in the corresponding indices in M the interpolation has linear complexity. For further details, please refer to Appendix A.5.3.
Grmping
k=
k= 13
k=
penquury perdhead
77.08
81.45
K3.55
79.03
82.71
83.71
por bunch
80.03
82.94
84.14
esoved porshaick
80.55
83.49
$4.19
[a] Per- Per-query
[b) Per head
(cl Per-batch Figure 4: Visualization of the
(d) Causal Per batch
M' E 1) 1 In the case of very long sequence, it is possible evaluates to and the subsequent in-
Group of each top-k method.
3.1.1 SPARSE ATTENTION AND FINAL OUTPUT CALCULATION WITHIN LINEAR COST
Sparse Attention Mechanism. We calculate re-weighted sparse attention probability A' by first applying a sparse masked matrix multiplication p. where p(Q, K MM) QK T c M followed by softmax operation a. Note than the Q and K matrices which are inputs to p are the same Q and K which were inputs to the Performer. We then re-scale the weights using Sprub E RT (defined
later in this paragraph), so that A Spmb @ @p(((, KT, M))) E R Note that is sparse operation and M' E {0, 1|7T is the previously obtained sparse binary mask using FlatCSR. The output of p only needs to store the non-zero values and indices, which are the indices where M* has value 1. The softmax o (p(Q, K" .)') calculates the softmax probability of non-zero values in the sparse input. After applying o, each row of the sparse attention matrix A' will sum to 1, therefore, due to the high sparsity of M the resulting non-zero values after a will be higher than the ground truth attention matrix from the teacher. To account for this effect, we scale the attention probability using a learned weight Sprb E RTT, where Sprob Spoob (Z) (Z was previously defined
5

---

Published as a conference paper at ICLR 2024
Memory
Latency
Validation Curve
70-
70
Durs (2=32)
150
Ours
Curs (kw64)
Perforer
60
60
Durs 1e-1201
Durs 16- 04)
100
Reformer
50
40-
Dora 15-1
Cur 2581
50
30-
E
30
19
Varila
Pertormer
Reformer
500
1000
1500
5
10 15
20
2000
4000
6000
10000
MB
(a) Computational costs vs. perplexity on Wikitex12 with OPT-125m. (b) OPT-125M validation curves
Figure 5: Fig. 5a shows how 9 variants of our model perform when comparing perplexity vs. computation. Baseline model marker sizes correspond to increasing numbers of buckets (Reformer) random projections (Performer). In all cases. SEA produces the best perplexity score. Fig. 5b Validation curves for SEA and baseline efficient attention methods. SEA converges much faster compared to Performer and Reformer.
Table 2: Comparison of perplexity score on Wikitext2 with various scales of OPT model. We trained the same number of steps (10k) for each method. We used k = 64; K = 64 for SEA OPT- 125M, 350M and 1.3B.
Perplesiit Mfter Adjusting k
Alurrs fromcedt
OPT-125M
OPT:350M
OPT:1.38
Method
PPL..
Mom.
Lat. L
PPI..
Mon.
J
tan. 1
PPI. |
Mem.
Lar. J
Vanilla
29.2
408
4.88
19.3
526
6.71
13.9
1120
16.49
Recommer
63.9 (+34.7)
902
10.90
58.2 (438.9)
1195
14.76
49.02 (-35 12)
2405
31.37
Performer
49.8 (+2204)
1.21
(417.)
60.5
1.82
30.6 (-167)
137
571
250
29
SEA (Ours)
26.0 (3.2)
6.76
19.5 (40.2)
241
9.43
(3.5
499
21.37
Figure 6: Dynamically ad-
justing after training.
in Section 3.1 CNN Decoder) and fprs is linear projection from Rd + IR followed by sigmoid activation function. Since the sparse calculation of p is only performed in the indices where M has 1, the complexity follows that of M which is (TT). The resulting A' matrix remains in sparse matrix format. Afterward, we calculate the context feature as C = A E RTxx which has linear complexity due to the linear complexity of A and V. The output C' is now stored in dense matrix format.
FlatCSR: A Modified Compressed Sparse Row (CSR) Format. We introduce our novel sparse operation, FlatCSR, which is an efficient implementation of the previously described sparse attention mask generation and attention operations utilizing grouped top-k selection. Our initial attempt for the sparse operations in our model utilized the Coordinate List (COO) sparse format. However, COO is not ideal because it stores every coordinate of each non-zero point and it does not take advantage of the structure provided by our grouped top-k selection, as shown in Table 3. Therefore, we eventually adopted the CSR sparse format instead of COO, as it uses less memory and schedules the per-row computations wisely using pre-constructed rows from our grouped top-k selection. We present detailed explanation of FlatCSR and further discussions in Appendix A.5.2.
Output Calculation. For the final output, instead of re- lying solely the sparse attention operation previously described, we combine the output of the Performer, and the sparse attention operation improve the ability of the highly sparse attention matrix to look up global information. The final output Cwsa is computed as a summation of two terms C and Cavg. each obtained from A and À respectively. First, we calculate the importance score of each token by averaging every row of Å, resulting in i = + En A E RRK We then interpolate to i E IRT and subse- quently perform weighted average pooling of Cavg = =iV E Rd. In causal attention, this global pooled context feature Cavg is replaced with an accumulated average of the tokens in V such that V, , E V, We mix Caug and C using learned weight 8mix Spool (Z) E IR.. with Spool composed of linear transformation and sigmoid activation fpoat : .Rd + R. Osea is calculated as C.e. Smix @ C + (1 m @ C avg
Table 3: Comparison of different sparse matrix formats on random inputs
Method
Latency (ms)
Memory (MH)
COO
75.66 (1002)
1191 (100%)
FlaCCS (Ours)
(15.06))
3.2 TRAINING SEA ATTENTION
For training SEA, we first replace the attention mechanism of pretrained teacher transformer model with our SEA attention mechanism. Then, we use knowledge distillation (KD) (Hinton et al., 2015) to train the newly added SEA attention parameters while adapting the original weights to SEA attention. Our training scheme is similar to previous transformer KD work (Jiao et al., 2020) since approximate the context features and attention matrices. However, we further add an objective

---

Published as a conference paper at ICLR 2024
Memory Map!
Hamory (WAPC)
Horoyyy (aLA
Namary (MNJ)
Accuracy After Adjusting R
x2
60
0.
2:44
(lagg
Litency UNIEes
Litency OHLLD
34
Sythaber
05 96
CA
06
04 0.5 80
85 2*
(b) Dynamically adjusting after training.
(a) Comparison of the tradc-off on GLUE
Figure 7: Fig. 7a shows the comparison of the trade-off hetween computational cost (latency and memory usage) and accuracy on each GLUE subset with BERT-base. 7a (column 1): Overall performance trade-off over three subsets. We average metries weighted with dataset size. 7a (column 2-4): Performance trade-off' per each GLUE subset. Baseline model marker sizes correspond to increasing numbers of buckets or random projections. Fig. 7b shows the accuracy trade-off of dynamically adjusting &s after training.
for the compressed estimated attention matrix A to match the teacher attention matrix. With C() signifying a loss for layer i, our overall training loss Lewa is given as the following, with each term described in the following paragraph:
Luea
1 T.
cl approx + prob Conneett
Lid
+ Lud tack + Lrask
To calculate Lapprox. we perform nearest neighbor interpolation to the estimated attention matrix A,, and get A RTXT in order match the shape of the teacher attention matrix A, E RTXT. Then we apply both KL divergence and an MSE loss between A! and A,, Cappoxx = (L A.) + MS A.). Next, we calculate </0 (A,,,,, Ā.) + MSB(A,, A.) which minimizes the error between the student's A; and teacher's A, attention matrices. For =prob we calculate the dense student attention A; = o (Q.K,T during training. We then add C MSE C() to minimize the error between the attention context feature (() and teacher context feature (+) E Rxxd Next, to minimize the error of each transformer layer (after the attention layer and MLP), we
c(i)
gather the outputs of each layer O sea,
(i) O NxTx xdd for SEA and the teacher, respectively, and
sea, (0)). The loss for training the i-th layer is Clu and is a weighted sum of each layerwise loss such that c(i) =
calculate c(u (SE(O (1)
((9) Capprox + prob + C + clG. We omit the weight term on each sub-task loss for simplicity; details are in Appendix 1.4.1, We then calculate knowledge distillation loss from the model output logits LraLms = KL(P, P), where P E R is model output logit. Finally, we sum together the average layer-wise loss and the downstream task loss Lusk into the SEA training loss given in Eq. (1).
4
EXPERIMENTS
4.1 CAUSAL LANGUAGE MODELING
We further evaluated SEA on the language modeling task on Wikitext2 (Merity et al., 2017) with various OPT (Zhang et al., 2022) variants, which involves causal attention. We selected two repre- sentative baselines, Reformer (Kitaev et al., 2020) and Performer (Choromanski et al., 2021) l). which represent the sparse attention and kernel-based linear attention methods, respectively. In Tables and A.8, Reformer shows unpredictable performance between the tasks, exhibiting strong perfor- mance in text classification (Table A.8), and the worst result in causal language modeling (Table 2). In contrast,
proposed method, SEA attention, performs the best in both cases with the closest perplexity score to the vanilla OPT and even surpasses the quadratic attention model on OPT- 12MM in Table 2. In Fig. 5a, we show trade-off between computational cost and perplexity. Our method exhibits more latency, since we utilize both hernel-based and sparse attention within our model (de- tailed lateney breakdown in Fig. 8). Therefore, we discuss the latency and memory trade-off of our method in Section 5. We note, however, that even though SEA uses both Performer and sparse at- tention modules, the convergence rate is much faster than both solo baselines, as depicted in Fig. 5b due to the direct attention distillation from the quadratic teacher.
4.2 TEXT CLASSIFICATION
We perform text classification evaluation of SEA on the GLUE (Wang et al., 2019) benchmark with BERT (Devlin et al., 2019). We train SEA attention by adapting the fine-tuned model, as
7

---

Published as a conference paper at ICLR 2024
Peak VRAM Usage
Latency
24
21
Fashvtsention
27
Costormer
MB
4
21
Performe
Reformer
Sinihom
2
Syntheszer
22
Oura (ied4]
24
2
IR/SR-1128
212
212
21
2
21
2
2
2
2
21
2
2
Sequence Length
Sequence Leigth
Performer (1608N)
Predictor 27 3990
Dense Operator
Mask (6.66)
Interp (A2P))
Sparse Attention por.2N0
SEQ-8192, e1iE K-256, 24 62m2
RetCSR Operator
Figure 8: (top) Space and time complexity comparison between SEA attention and baselines. Lower values are better in both figures. SEA exhibits complexity in line with other linear attention models. We show a trend- line for the quadratic attention because it runs out of memory on sequence lengths larger than 2 (bottom) Latency (ms) breakdown of our SEA attention (1" = 2 cousal-per-batehl Each orange and green box shows that the operation is computed in our novel FlatCSR format and in dense tensor format, respectively.
described in Section 3.2. In Fig. 7a, we show a trade-off between computational costs and various performance scores. We test the following baseline methods: Reformer, Sinkhorn (Tay et al., 2020b), Performer, Cosformer (Qin et al., 2022), and Synthesizer (Tay al., 2020a) (see Appendix A.4.2 for further experiment details). In all tested subsets, SEA achieves the top performance and maintains competitive latency and memory costs. In Table A.8, we show results of the tested baselines within the same constraints by limiting all the baselines and our method to have the same bucket size in sparse attentions and the same number of random projection feature sizes in kernel-based attentions.
To summarize, method achieves higher accuracy than all linear baselines while maintaining competitive performance in terms of latency and memory usage. SEA comes within 0. 0.1% accuracy of quadratic attention on MNLI in Table A.8, and in Section 4.3 and Figs. 6 and 7b we show we can dynamically adjust after training to outperform quadratic attention.
4.3 DYNAMICALLY ADJUSTING k
In Figs. 6 and 7b, we experiment with dynamically adjusting l after training with fixed value of l on the Wikitext2 and MNL. dataset. We find that increasing also increases the accuracy without the need for any further training. This means even after fine-tuning the SEA, our model still preserves pretrained knowledge and increases accuracy when the constraint on l is relaxed. Therefore, this characteristic helps users to design flexible and dynamic models that can adapt to real-time service demands and cost constraints by dynamically adjusting At For example, considering that lower leads to lower computational cost shown in Fig. 5a and Fig. 7a, if given situation calls for lower computational cost, A: can be minimized, while if accuracy is more important, k can be set to higher in real-time. In addition, surprisingly, increasing tt after training makes the model perform better than the vanilla quadratic model. In Fig. 6, the vanilla baseline shows a perplexity score of 29.2, however, all SEA models (k- 32, 64, 128) surpass this when we increase k after training.
5 EFFICIENCY OF SEA ATTENTION COMPUTATION
In this section, we provide the memory usage and lateney experiment results of our method with different sequence lengths T. In Fig. 8, we show that our resource usage tendency is O(T). We test SEA attention with the causal per-batch top-k grouping mode with our FlatCSR implementation.
Peak Memory Usage. In Table A.I and Fig. (top-left), compare peak memory usage in mega-bytes for each attention method. Compared to baselines, SEA attention shows competitive peak memory usage. Our method shows an 81.05% reduction in peak memory usage compared to quadratic attention at sequence length 23 Our method consumes memory only about 78.81 % compared to Reformer, while consistently maintaining higher accuracy as shown in Figs. 5a and 7a and Tables 2 and A.8. Moreover, our methods successfully operate with competitive memory budget with other linear attention methods on all sequence lengths (shown in Table A. and Fig. 8), while quadratic attention exceeds memory capacity above 2 In summary, memory complexity to (T)), and exhibits less memory usage than Reformer.
method reduces
Latency. In Fig. 8 (top-right), we compare the latency between SEA and our linear attention baselines, showing that SEA scales linearly. Our model only incurs 32.72% of the latency cost of quadratic attention in Fig. 8 for sequence length of 213 where quadratic attention runs out of
8

---

Published as a conference paper at ICLR 2024
Teacher Adentien
Teacher Attention Matris 0274
Estinoned egg1
Compressed Estimated
Spersa Ineenton
Attention Matrix 02/1
MNL/with BERT
Wkitse22 OPT-125m
Figure 9: (Ieft) Intermediate attention examples. (right) The first row the attention probability of the teacher model, the second row is the compressed attention interpolated to full size. Interpolation the full size attention matrix is for visualizing estimated attention Á and is not part of the regular linear inference procedure. (a) MNLI with BERT-base (K" 128) (b) Wikitex12 with OPT-125m (K" 256).
memory. SEA also shows better performance with a similar latency to Reformer, as shown in Fig. 7a (bottom-leftt. However, our method also shows a latency-accuracy trade-off in Fig. 7a, where some baselines such as Sinkhorn, Cosformer, and Performer show better latency but worse accuracy than our SEA. We break down the latency of each component of our proposed method in Fig. 8 (bottom). The dense operations mainly permute and reshape, comprise 6.27% of latency. However, in the COO sparse format, the dense operations use 13.31 %. and COO sparse operations comprise 86.68% of the latency cost. As a result, the COO format is 6.63x slower than our novel FlaCSSR format as shown in Table 3.
47.45%, FlatCSR sparse operations use 46.28% and the other operations,
VISUALIZATION OF ESTIMATED ATTENTION FROM SEA ATTENTION
In Fig. 9 (right-a), using BERT and the MNLI dataset, we visualize the interpolated estimated at- tention matrix Á from SEA and compare with the attention matrix of the teacher model Á. The learned estimator of SEA attention shows the ability to pre- dict various attention shapes from the original fine- tuned BERT model. As be seen, estimator learns well-known fixed patterns, such as the diag- onal but also dynamic patterns that require contextual interpretation. In Fig. 9 (right-b), we show the visualization of causal attention commonly used in generative models. In the causal-attention setting, we observe a diagonal attention probability with wavy or chunked diagonal lines, patterns that cannot be handled by previous heuristic linear attention mask patterns. However, our estimator still shows great predictions on such highly variant patterns. In addition, in Fig. 10, we show our compressed attention Á, top-k compressed mask M, sparse mask M and sparse attention A".
Atertion
Mack
Sparse Vank
Sporss Amanntion
M
A
Figure 10: Visualization intermediate buffers during masking and sparse attention.
Moreover, our model can perform well even if the estimated attention is slightly different from the teacher's, thanks grouped top-k; which drops values that are not selected the top-k pro- cedure. For example, in Fig. 9 (left-bottom), we show sparse attention matrix after masking the estimated matrix with our grouped top-k selection masks. Although the estimated attention matrix seems somewhat different from the teacher's Fig. 9 (left-middle), the resulting sparse attention pat- tern Fig. 9 (bottom-left seems quite similar to the teacher's after applying the top-k mask. Further visualization results can be found in Appendix A.2.
7 CONCLUSION AND DISCUSSION
Our proposed method, SEA attention, shows state-of-the-art performance for integrating linear tention with pretrained transformers, as we show empirically in Section 4. The critical change over existing works is that we estimate the attention matrix compressed size using dernel-based lin- ear attention to form compressed sparse attention mask which can be decompressed into a full sparse attention mask to overcome the quadratic cost. By doing so, we can preserve the dynamic and complex attention patterns of the pretrained teacher transformer model through direct attention matrix distillation. Furthermore, SEA also provides interpretable attention patterns. SEA performs similarly to vanilla attention while existing works could not. We look forward to seeing future search that; may apply a learnable masking method instead of top-k selection, such as concrete masking (Lee et al., 2023), or improve our uniform interpolation by some non con--nfform or learnable interpolation which may provide further performance increases.

---

Published as a conference paper at ICLR 2024
ACKNOWLED
We extend our heartfelt appreciation to Seanie Lee and Minki Kang for their insightful reviews, which greatly enriched the quality of our work. This work was supported by the National Re- search Foundation of Korea(NRF) grant funded by the Korea governmen((MSIT) (No. RS-2023- 00256259).
REPRODUCIBILITY STATEMENT
We will introduce code construction and data collection in this section for reproducibility.
SEA Attention Module Implementation First of all, we use perlin as our code name of SEA attention on supplementary source code. We implement SEA attention modules, including self-attention and causal attention for the transformer encoder and decoder. Users can import sre. brs.models per in.attention module to construct custom transformer model. We imple- mented sample transformer models for experiment SEA attention in this paper: BERT and OPT. Users can check SEA-BERT implementation in sre. models per_in.eert, and SEA-OPT im- plementation in sIC, mwdels.do.....o pww.in.oop......
FlatCSR Triton Implementations We implemented our customized FlatCSR datatype operations in sre wwwels wtw.lls.pprrin... wwwppeeeatteetiiiiii... ops.. Each kernel definition file has its own benchmark logic and verification logic.
Training Please check README.md in the root directory of the supplementary source code archive. We provide detailed guides and Python and Anaconda environment files to repro- duce our results. Users can easily test various configurations prcipts/opt... script. Users can easily test various configurations of SEA attention BERT with src. trainer. wwt.:pperrnnnninneneeee ule, users must supply teacher attention scores and the context layer of attention layers to cal- culate the SEA attention loss. We implemented examples of those pipeline in are.trainer.w, aro.models hh..ppt....t....p and src models. hf bert.
SEA attention on OPT with
program entry. To train custom SEA attention mod-
Testing Please check README.md and sze tests...sss....... We implement several test codes to validate our implementation. For example, verifying the causality of causal attention.
REFERENCES
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longtormer: The long-document transformer. CoRR, abs/2004.051555 055500 2020. URL httph:://///aeeee org/ abs /2004 05150. 1,
Daniel Bolya, Cheng- Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman. Token merging: Your vit but faster. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL /tppr//ieeeeeeeieeee net wttpppeeeieeeeeeeeeeeeeeeeetttte httpp//dddddooaarrrrT
2
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, 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 McCan- dlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tiee Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https: httt:p:://prooeedinggss... 1457c0d6bfcbe 962204967800 14 html.
mew.:eeeeeddd........e..ee.eeeee neurips cc wppppa//a///////
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song. Atri Rudra, and Christopher Ré. Scatter- brain: Unifying sparse and low-rank attention. In Mare' Aurelio Ranzato, Alina Beygelz- imer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances
10

---

Published as a conference paper at ICLR 2024
in Neural Information Processing Systems 34: Annual Conference Neural Informa- tion Processing Systems 2021, NeurlPS 2021, December 6-14, 2021. virtual, pp. 17413- 17426. 2021. URL https: http://poeeeddd.nerrreeeeeeeeeee
CC wwtpp/eppppppepceeeeeeeee html. 1, 1,
Ruijun Chen, Jin Wang. Liang-Chih Yu, and Xuejie Zhang. Learning to memorize entailment and discourse relations for persona-coneee consistent dialogues. In Brian Williams, Yiling Chen, and Jennifer Neville (eds.). Thirty-Seventh AAAI Conference on Artificial Intelligence, AAAI 2023, Thirty-Fifih Conference on Innovative Applications of Artificial Intelligence, IAAI 2023, Thirteenth Sympo- sium on Educational Advances in Artificial Intelligence, EAAI 2023, Washington, DC, USA, February 7-14, 2023, pp. 12653-12661. AAAI Press, 2023, doi: URL https: /tt:://ooiiooo org 1.. 1099 www11660/a/a...... www.w1. 26489.
v37111.26688)
Wei-Lin Chiang, Zhuohan Li. Zi Lin, Ying Sheng. Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, lon Stoica, and Eric P. Xing. Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023. URL https:
/ http://ggggbb/ggo////////
03 30 vicuna/.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quiney Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with per- formers. In 9th International Conference on Learning Representationss ICLR 2021, Virtual Event, Austria, May 3-7, 2021. Open 2021. URLhttps: /wwpr/eeeeeeeeeeeeeeeeee...eeeee net forum? id=Ua6zuk OWRH. 1, 2, 3, 4, 7, 14
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359, 2022. 14, 20, 22
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, NATCL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Pupers), pp. 4171- 4186. Association for Computational Linguistics. 2019. doi: 14288663/111111-14433. URL 2, 7, 14, 15
/oii. www./oo../....ooooo.ooooo httg://..115588 18e111863//// wwt.wlwwwilliii1 https: /doi
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x 16 words: Transformers for image recognition scale. In 9th International Conference on Learning Representutions, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. WRRLLttps://openzeeet id-YichE TNTTy. 1, 2, 20
net forum?
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus, 2019, URL http: htt///yyloooo0// .i gii.ttttiii tit.......... i WpwneeeeeTettCorprtruu 21
Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in a neural network.
CoRR, abs/1503.0255 2015. URL http: / arxiv org/ abs / 1101. 02531. 2, 6
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling BERT for natural language understanding. In Trevor Cohn, Yulan He, and Yang Liu (eds.), Findings of the Association for Computational Linguistics: EMNLP 2020, Online Event, 16-20 November 2020, volume EMNLP 2020 of Findings of ACL, pp. 4163-4174. Association for Computational Linguistics, 2020. doi:
URL httphttps:/ddde org/ 0ee.g//..111.13. 0...8633v1111 /2020 f findings wwwnddinnnnmllllnm 372. 2, 6
Schoon Kim, Sheng Shen, David Thorsley, Amir Gholami, Woosuk Kwon, Joseph Hassoun, and Kurt Keutzer. Learned token pruning for transformers. In Aidong Zhang and Huzefa Rangwala (eds.), KDD '22: The 28th ACM SIGKDD Conference Knowledge Discovery and Data Mining, Washington, DC, USA, August 14 18, 2022, pp. 784-7944 ACM, 2022, doi: 10. 145/5447333678. 3539260. URL https: /do. org 11 14 3534 3533363396 2

---

Published as a conference paper at ICLR 2024
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL https: htttps://ppnreriew..ee net.:iiwwwweeeeeeeeeeee.ee.eeeee www.:nnnne/ototttttoooooo.
rkgNKkHtvR. 3, 7, 14
Heejun Lee, Minki Kang, Youngwan Lee, and Sung Ju Hwang. Sparse token transformer with attention back tracking. International Conference on Learning Representations, 2023. URL
https: /ttp::///ppeeevviiee net/ forum 2,9
Liu Liu, Zheng Qu, Zhaodong Chen, Yufei Ding, and Yuan Xic. Transformer acceleration with dynamic sparse attention. CoRR, abs/2110.1 1299, 2021. URL
2110.11291 1,3
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. https: / /popereeview tatt/forrmmm id=By j 72uxxe. 2, 7, 16
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenRe- 1,2, 8,
view.net, 2022. URL WRLhttps: http::::////////vvvveee net /tt///oommmmmmmmeeeee0 14
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention in transformer models. CoRR, abs/2005.00743, 2020a. URL https:
alaxiiv.org abs 2005. 00743. 3, 8, 14
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 9438 9447. PMLR, 2020b. URL http: http:///occeeddeeeee . mlr. wptps//eie/ttaaeeeeeeeee 14
html. 1,3,8, 8,
Philippe Tillet, Isiang-Tsung Kung, and David D. Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Tim Mattson, Abdullah Muzahid, and Ar- mando Solar-Lezama (eds.), Proceedings of the 3rd ACM SIGPLAN International Workshop Machine Learning and Programming Languages, MAPL@ PLDI 2019, Phoenix, AZ, USA, URL https:
June 22, 2019, pp. 10-19. ACM, 2019. doi: 10. 00.1145/3115008.322997 //doi. org.orggggoo.ooo. 10. 1141 /1/44//33555 3329973. 18
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton-Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao. Vedanuj Goswami. Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura. Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee. Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet. Todor Mihaylov. Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Hllen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurélien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine- tuned chat models. CoRR, Abs/2307.09288. 2023. doi: 10.48855/anxii..337..222 09288.. URL https: /ooi. www.iiorrrl..... 48550/ arxiv. 2307 09288.
Aäron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew W. Senior, and Koray Kavukcuoglu. Wavenet: A generative model for raw audio. In The 9th ISCA Speech Synthesis Workshop, Sunnyvale, CA, USA, 13-15 Septem- ber 2016, pp. 125. ISCA, 2016. URL http:/ / /wwww www.:iiscasseeleeee wwweaaapppec........ooooooo. www.gg//rrvvvveeee SSW 201 w/t.st/aattss///eeeosss ssw 9 DS - van den Cord. html. 18
12

---

Published as a conference paper at ICLR 2024
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp- 2017. URL https: http:////ppcceeeeeeeeeeee new.iiiiinnnnnnn.s...........iss wwtprripppccpppppppeeeeeee http//aapr///11////eee 20th//ahh/h//// 3f5ee243547dee wwd.033lcc488555e hat.e22223eeeeooooccccccoooooo..
Abstract html.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OvenRevie 2019. URL https: httpp::///pemeee http///pppnweew.eee/ee/ee httpriiwwweeeeeeeeeee net/forumiid rJ4km2R5t7. 2, 22
Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaei, Atharva Naik, Arjun Ashok, Arut Selvan Dhanasekaran, Anjana Arunkumar, David Stap. Eshaan Pathak, Giannis Karamanolakis, Haizhi Gary Lai, Ishan Purohit, Ishani Mondal, Jacob Ander- son, Kirby Kuznia, Krima Doshi, Kuntal Kumar Pal, Maitreya Patel, Mehrad Moradshahi, Mihir Parmar, Mirali Purohit, Neeraj Varshney, Phani Rohitha Kaza, Pulkit Verma, Ravsehaj Singh Puri, Rushang Karia, Savan Doshi, Shailaja Keyur Sampat, Siddhartha Mishra, Sujan Reddy A, Sumanta Patro, Tanay Dixit, and Xudong Shen. Super-naturalinstruccii declarative instructions on 1600+ NL.P tasks. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.). Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022, pp. 5085- 5109. Association for Computational Linguistics, 2022. doi: 10.18663/11202222mmpmmmmmm URL https: /ooi http://oo...o111111151133311/
Generalization via
186118666333 2022 emnlp p- 340.
Adina Williams, Nikita Nangia, and Samuel R. Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Marilyn A. Walker, Heng Ji, and Amanda Stent (eds.), Proceedings of the 2018 Conference of the North American Chapter of the As- sociation for Computational Linguistics: Human Language Technologies. NAACL-HLT 2018, New Orleans, Louisiana, USA, June 1-6, 2018, Volume I (Long Papers), pp. 1112-1122. As- sociation for Computational Linguistics, 2018 doi: 10.18653/11n18--11
URL https:
//doi. www.ooo//.......... (188886//lv/nn - 2. 15
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Trans- formers for longer sequences. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, De- cember 6-12, 2020, virtual, 2020. URL ht httppss:/ppcccccinnnttet/ htwpeeeeeeiinnn...uii.......e.ee catsscccppaae cc/paper////ppccpp/coo.cc..eeepp 2ttpa//.....22222883332213111111
attsaaaaaabbbbbbaaaattttaaaaaate htm1. 1,3
Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama, Sep 2023. URL https:
//github.r com wtt. jttpnn//i/ngg///// / ama. 21
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christo- pher Dewan, Mona T. Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettle- moyer. OPT: open pre-trained transformer language models. CoRR, abs/2205.01068 2022. doi: 10.48550//rXiv.2205 2.00002200..0008. 01068. URL ht https: : / //w.i//iio////////.../ org/10. 48550/ www.o//a.ivvv... 2205 01068. 2, 7, 14, 16, 18
13

---

Published as a conference paper at ICLR 2024
A APPENDIX
A.1 EFFICIENCY MEASURES OF SEA ATTENTION
Table A.I: VRAM usage comparison between baseline and SEA attention. The unit is MB, lower is better.
Fach column represents different token length setting random inputs.
4,096
8,192
16,384
32,768
65,536
131,072
Avg.
Vanilla
134.00
524.00
2072.00
OOM
OOM
OOM
OOM
FlashAttention
6.00
12.00
24.00
48.00
96.00
192.00
63.00
Performer
11.66
23.29
46.54
93.04
186.04
372.04
122.10
Sinkhorn
20.03
40.09
80.32
161.15
324.26
656.50
213.73
Cosformer
28.14
56.16
112.19
224.25
448.47
897.12
294.39
Reformer
218.65
469.26
1066.72
2645.03
7338.06
OOM
OOM
Synthesizer
134.02
524.04
2072.08
OOM
OOM
OOM
OOM
Ours (k=32)
48.64
97.35
195.00
391.19
786.86
1684.79
533.97
Ours (k=64)
49.75
99.39
198.76
397.55
795.14
1684.79
537.56
Ours (k=128)
54.84
109.47
218.77
437.55
875.08
1751.71
574.57
Table A.2:
Latency comparison between baseline and SEA attention. The unit is milliseconds per iteration, lower better. Each column represents different token length setting random inputs.
4,096
8,192
16,384
32,768
65,536
131,072
Avg.
Vanilla
1.66
8.61
35.66
OOM
OOM
OOM
OOM
Performer
0.25
0.49
0.90
2.00
4.25
10.94
3.14
Cosformer
0.36
0.73
1.49
2.94
6.28
12.82
4.10
Sinkhorn
0.21
0.44
1.00
2.41
6.42
18.68
4.86
FlashAttention
1.30
5.17
20.57
81.73
330.19
1348.04
297.83
Reformer
2.33
5.26
12.24
29.94
83.41
OOM
OOM
Synthesizer
1.53
8.04
32.42
OOM
OOM
OOM
OOM
Ours (k=32)
1.45
3.07
6.21
12.00
25.74
50.89
16.56
Ours (k=64)
1.77
3.69
7.40
13.46
28.74
56.46
18.59
Ours (k=128)
2.39
4.89
9.81
16.49
37.94
72.59
24.02
In this section, we show detailed results of efficiency measurements from SEA attention. We tested various sequence lengths with various attention methods: none (Vanilla), Sinkhorn (Tay et al., 2020b), Cosformer (Qin et al., 2022), Performer (Choromanski et al., 2021), Reformer(Kitaev al., 2020), FlashAutention (Dao et al., 2022) and Synthesizer (Tay et al., 2020a). The test is performed in the bidirectional self-attention setting. We K 128 for SEA attention. We show memory usages in Table A. and we show latencies in Table A.2. We execute all benchmarks on the same machine with the same resources. Our environment is built with Ryzen 3950x, RTX 2080ti 8x PCIe 3.0, DDR4-2400 64GB, and Ubuntu 22.04. The versions of third-party libraries including PyTorch and Triton are described the supplementary file, requirement txt. Also, we pro- vide the docker environment of our experiment environment for reproducing results, done with the supplementary file, DockerFile.
A.2 ESTIMATED ATTENTION VISUALIZATION OF SEA ATTENTION
We show a more detailed attention estimation visualization in Figs. A.1 and A.2. We visualize teacher ground truth attention, estimated attention, student attention before masking, and student sparse attention after masking for each layer and head of BERT-base (Devlin et al., 2019) and OPT- 125m (Zhang et al., 2022).
14

---

Published as a conference paper at ICLR 2024
(a) Attention matrices of the first 6 layers
Attention matrices of the last 6 layers
Figure A.I: Visualization of intermediate attention matrix buffers in SEA attention of BERT-base (Devlin al., 2019) on MNLI (Williams et al., 2018). visualize teacher ground truth attention, estimated attention, student attention before masking, and sparse student attention for each layer and head, with k=13. Each group of rows shows attention matrices from one layer. We stack the visualization layer by layer vertically, from top bottom, while showing all heads in each particular layer horizontally, from left to right. We show the first layers in Fig. A. and last 6 layers in Fig. A. Best viewed with high zoom.
15

---

Published as a conference paper at ICLR 2024
bod
Cmaed
Erand
Jnamosn
Huter Jabar
Teahs
Adhetion
Talhs
Infer Mesiing
Talhs
Bulmadad
(a) Attention matrices first 6 layers
(b) Attention matrices the 6 layers
Figure A.2: Visualization intermediate attention matrix buffers in SEA attention OPT-125 (Zhang et al., 2022) on Wikitext2 (Merity et al., 2017). We visualize intermediate buffer samples from attention matrices here, the same way as Fig. A.1. We show the first 6 layers in Fig. A.2a and last 6 layers in Fig. A.2b. Best viewed with high zoom.
16

---

Published as a conference paper at ICLR 2024
A.3
VISUALIZATION THE MASKING PROCESS
Estimattdd
Estimated
Sparse
Sparse Attention
Sparse Interpolation
Top-k Selection
Sparse Attention the Mask
Figure A.3: Visualization of masking process intermediate buffers during attention estimation in compressed size and sparse interpolation performing sparse attention.
In Fig. A.3, we visualize the intermediate buffers for masking sparse attention using M for better understanding. The example is sampled from OPT-125M, which uses causal attention. We show the process to perform sparse attention using the mask M estimated with compressed attention matrix estimation Á. the visualization, we differentiate binary masks and real buffers by using black- and-white and red-and-black color schemes. For note, the sparse matrices M and A converted into dense matrices format in order to render the image. Black represents zero-valued pixels, which stored in memory. Since the visualized attention mechanism causal attention, each row of the compressed estimation Á and M is resized with different target widths according to the token index in M and A'.
A.4 EXPERIMENT DETAILS
A.4.1 TRAINING HYPERPARAMTERS
Dataset
MNLI COLA
MRPC
Wikitext2
Batch Size
32
32
Table A.3: Batch sizes for various datasets in our experiments shown in Scction 4
Loss
KLL Caypex
Lypaox
Lyne
Cyuis
Loontest
Cya
Ess usk
Cus
Weight
0.1
0.1
1.0
5.0
0.1
Table A.4: Loss weights for different loss terms defined in Section 3.2
Batch sizes for experiments outlined in Section 4, be seen in Table A.3, We define different learning rate values for original parameters and SEA attention parameters. We use learning rate 10-5 for original parameter, and 10-4 for SEA attention parameters. For OPT models, learning rate * 10-6 for the original parameter and 10-4 for SEA attention parameters. Weights for loss scaling outlined in Section 3.2 can be seen Table A.4,
A.4.2 GLUE
We test SEA attention with settings à {7, 13, 13,25,,, 25} and K {32, 64, 128]. We changed the bucket size to match the sparsity constraint in Reformer and Sinkhorn and the number of base projection feature sizes in Performer. We test attention methods within the fixed sequence length (256) to
17

---

Published as a conference paper at ICLR 2024
measure latency and memory usage. We train all methods, 20 epochs in MNLI and 50 epochs in COLA and MRPC.
A.5 IMPLEMENTATION DETAILS
A.5.1 ATTENTION ESTIMATOR CNN
Before arriving at the attention estimator CNN. there are two MLP's μ. RM + Rd and v Rd H Rd R Rea//e. which projects the kernel-based attention output such that 2 v(u(V BDxTxKea/ss This is then transposed and resized to be E R Me, * XTxK/e, as explained in Sec- tion 3.1. v expands the channel dimension to and reduces the width hidden state Empirically, we found that the channel expansion helps the CNN learn a better encoding, and the size reduction reduces the overall computation cost of the CNN. In our experiments, we set C. = 2 and e/ = 4. After obtaining Z, we decode it using the 3-layer CNN. The first convolution layer reduces height by Co using kernel size 3; s RR xT X/cc >> HR xT/,,,xT/,xxx/// The second layer performs another convolution using kernel size of 3: (2) IR NT/.,xxX//ss x K/e, + le xT x k/e, Then we resize the hidden state using the nearest neighbor interpolation to make the output R
H. xTxK R
The last layer changes the channel into the number of heads; /102 RR=
Lastly, we perform a softmax operation, finally obtaining Å. In causal attention, we do not re- duce the height. We only reduce the width to reduce computation. If one needs deeper CNN, then the second layer can be duplicated multiple times. Additionally, when applying SEA on large scale pretrained language models such as OPT (Zhang et al., 2022), fr accepts token embedding H-dd instead of the single head embedding R&dd so that it may learn information across the large group of attention heads. For causal attention, we change Vi into learnable positional embedding and use causal CNN (van den Oord et al., 2016) to satisfy to the causality condition, see Ap- pendix 1 for details. The implementation of the CNN can be found in the supplementary file at: sre. modelolll....m per peerrinnnaaeee innaateetiiinnnnnn attention. wew.eerntteeetio....
A.5.2 FLATCSR: MODIFIED CSR FORMAT To HANDLE GROUPED MASK
Here, we provide a detailed explanation of our novel sparse operation, FlatCSR. Our initial attempt for the sparse operations in our model utilized the Coordinate List (COO) sparse format. However, COO is not ideal, as storing full coordinates for every point in a sparse matrix makes the per-row computations difficult since each row must be identified and constructed from raw coordinates, as shown in Table 3. Therefore, we eventually adopted the CSR sparse format instead of COO, as it uses less memory and schedules the per-row computations wisely using pre-constructed rows. However, when it comes to causal-per-butch. which is our recommended setting in most cases, there exists challenge with the non-contiguous top-k grouping in the attention matrix since requires flattening the head and query dimensions. Therefore, FlatCSR, utilizing the Triton compiler which compiles Python code into low-level CUDA kernel binary (Tillet al., 2019). FlatCSR is capable of handling non-contiguous flattened tasks within the GPU kernel. In this paper, we implement FlatCSR only for causal per-batch. We note, however, that we expect the same number of non-zero entries in each of the top-k strategies depicted in Fig. 4, and therefore all strategies should show the same memory usage and latency.
propose specialized CSR tensor operation, called
We implement the interpolation and attention operations for the causal-per-batch grouping described in Section 3.1 in a sparse CSR tensor by transposing the head and query dimensions and flattening the interpolated attention matrix's 's last two dimensions (the head and key dimension). Ideally, we can store our attention mask with the CSR tensor because we have a similar number of non-zero entries per row (query) and the same number of non-zero entries per batch. We name this CSR tensor of transposed and flattened attention mask the FlatCSR tensor in this paper. However, to use this FlaICSR tensor in linear algebra operations, we must reshape and transpose the CSR tensor. Therefore, we implement new GPU kernel that internally performs reshape and transposes from the FlatCSR using Triton (Tillet et al., 2019). We heavily utilize the property that every row (query) has similar number of non-zero entries (approximated scheduling. Therefore, we can be more efficient in terms of memory and computation than the COO tensor type, which is generally used in sparse tensor computation.
k) during memory allocation and thread
18

---

Published as a conference paper at ICLR 2024
A.5.3 SPARSE INTERPOLATION
In this section, we describe the details of sparse interpolation from M M We interpolate M {0, 1)T*K into M E {0, 1)TxT as described in Section 3.1. We claimed that the complexity of this interpolation is O However, if we perform the interpolation of the M in a dense matrix for- mat, the complexity should be (TT2). Since we perform the interpolation in a sparse matrix format, we only need to calculate the interpolation of non-zero entries in M. This is possible because we interpolate binary masks using nearest-neighhor as there is no requirement for linear or non-linear interpolation between pixels. 'he nearest neighbor interpolation is independent of other nearby pixel values and only depends on pixel indices, which are stored the sparse matrix format. This allows us to perform interpolation within (number of non-zero entries after interpolation complexity. We adjust the sparsity of M (which is determined by k) to make M have a constant number of non-zero entries (T * k). As a result, we always know how many pixels are in M and that the number of pixels OTT). In summary, the only thing we need to do interpolation is iterate every non-zero pixel in M and duplicate or reduce the number of pixels which are output to M depending on the pixel location in M and the ratio between T, K.
M")
However, to keep the linear complexity of sparse interpolation need to deal with the case that kK < T in other words kK/T < 1. This case is important because if this occurs, then k = max(1, round * K/T)) will be evaluated to since we always must select at least one pixel from A. If we interpolate the M in that case, then each pixel replication of the non-zero entries in M should be limited to k. The reason is that we set the lower bound of k into 1 to avoid an empty attention mask when à is zero after rounding because T is much larger than K. However, in that case, single pixel in M will be [T/K] after replication, and the total number of non-zero entries in M will be quadratic. We can solve the problem by limiting the upper bound of pixel replications to k* = min (k: [T/ because the k is always 1. However in this case we will encounter how we select k* pixels among [T/K]] pixels because the originally selected pixel in M covers [T/K] pixels in M In this paper, we decided to sample uniformly.
However, we think uniformly sampling the attention relations among block regions is not ideal be- cause there a high chance of not selecting higher probability attention relations because a uniform sample does not consider the distribution of attention probability. We would be interested in see- ing future research that deals with this case, which is what we expect to happen more often in the future LLMs keep inereasing their context length (OPT (2022)'s context length: 2048, LLaMA2 (2023)'s ' context length: 4096).
A.6 FLOPS COMPARISON
FLOPs Comparision Between SEA (Durs) And Quadratic Attention
350
Quadrati:
Ourstk:
300
Oursik-64 K-64)
0urs/c-128, K-G4)
200
50
2000
4000
6000
14000 Sequencee Length
8000
10000
12000
16000
Figure A.4: FLOPs comparison with SEA and quadratic attention.
We compute the FLOPs of our SEA attention and quadratic attention with different hyperparame- ters (k 32, 64, 128, K 128) in Fig. A.4. SEA shows clear linear computational complexity. However, the original attention mechanism shows a drastic increase in computational cost on longer context length due to quadratic complexity. The figure shows only FLOPPs and arithmetic operations.
19

---

Published as a conference paper at ICLR 2024
Therefore, this tendency is the ideal latency complexity of our method and means that there is plenty of room for improvement in implementation.
A.7 DISCUSSION ABOUT DIFFERENCE BETWEEN SEA AND FLASHATTENTION
In Fig. 8. we show two quadratic attention baselines, vanilla, and FlashAttention (Dao et al., 2022). While FlashAttention consumes memory linearly, it scales quadratically in terms of computation complexity. FlashAttention is an efficient implementation of original quadratic attention that elim- inates most of the memory space and bandwidth requirement of the attention mechanism by fusing attention probability calculation and context vector calculation. FlashAttention has linear memory complexity and quadratic time complexity because does not require space for storing attention score matrix to compute softmax probability. The lack of attention probability storage is an advan- tage in terms of memory bandwidth consumption. However, this characteristic may be a downside when attention probability is required in usage scenarios where we are concerned, such as token importance analysis for token compression. Therefore, FlashAttention has some of the same limita- tions as previous linear attention methods. However, our SEA attention does not have such limita- tions by providing an estimated attention matrix, while showing linear complexity in both memory and computation.
A.8 EVALUATION OF SEA ON LONGER CONTEXT MODEL
Table A.S: Evaluation result of context length extension from 2048 4096 on WikiText2. The reformer was not included in this experiment due to the already poor performance seen in Table
OPT-125M
Method
PPL. (after interp.) 4
PPL. (trained) & Memory (MB) 4
Latency (ms) |
Vanilla
52.96
18.96
1584
17.86
Performer
68.29 (+15.33)
62.40 (443.44)
102.75 (6.48%)
2.38 (13.32%)
SEA (Ours)
43.96 (-9.00)
23.43 (44.47)
448 (28.28%)
14.31 (80.12%)
We evaluate our method with longer context on Wikitext2. However, OPT only supports context lengths up to 2048, therefore we used positional embedding interpolation that was introduced in previous work (Dosovitskiy al., 2021), using bilinear interpolation. After interpolating the posi- tional embedding, we perform a few optimization steps (750 steps. 1.5M tokens) on the model with the causal language modeling task loss. In longer context length, 4096, our method outperforms quadratic attention in both latency and memory cost. Also, the experiment result shows our model is much stronger than baselines after positional embedding interpolation. This result interesting, as we think this shows that sparse attention masking helps to preserve the important attention relations by masking out non-important attention relations. We picked the SEA-OPT result from Table A.5, and details are following.
Table A.6: The trade off on long context (T = 4096) experiment on Wikitex12 using post training compression techniques: Query skipping and dynamic control. Each entry of table shows PPL (ms/ /MB). Each values colored with green and red. Better values are more green, and worse values are more red.
Opack &
112
128
Quer Wige
16 23:43 (14.31 /448.37)
23,00 (55.70/455.66)
22.63 (15,40 463.06)
22.34 (16.06 4770.28)
22.11 (16.65 (777.59)
8
23,26 (14.61/444) 4448.37)
22.82 ((8.2884455.65
22.48 (18774/446306)
22.20 (16.35 470.28)
21.98 (16.97/ /477.99)
4
23.02 (15.33 (448.7)
22.60 (15.65/455.66)
22.29 (16.467463.00)
22.02 (47.03/:47028)
21.83 (17.59 477.899
2 22.73 (16.35 (16.354448.57 /448.37)
22.33 (16.86/4555665
22.02 (17.50 463.06)
21.78 (18.08 / 470.28)
21.58 (18.64/ 477.58)
1 | 22.38 (18.17444.35)
22.04 (1868/455.66)
21.76 (19.31 462.96)
21.55 (19.92/470.19)
21.38 (20.477/477.54)
In Table A.5, we show the performance trade of the longer context model using our post-training compression techniques: query skipping and dynamic k control. We newly introduce the query skipping in this section. Query skipping is skipping rows before CNN inputs, and replicating rows after CNN, reducing the cost of the CNN decoder. We trained the OPT-125m with self distillation setting because we do not have a properly trained 4k OPT model, therefore we use itself as a KD)
20

---

Published as a conference paper at ICLR 2024
teacher. Also in this experiment use A 128, K 96 to train efficiently longer sequence length. On each transformer layer, we use imput Q. K as a source of the teacher attention matrix. For computing self teacher attention matrix. we cut the gradient of each Q, K to prevent training from oscillating and exploding due to self feedback loop of the gradient. Table A.5, the model shows better latency and less memory consumption toward top-left, and gets closer to original model performance toward to bottom-right.
A.9 EVALUATION OF SEA ON LARGER DATASET
Table A.7: Evaluation result on DeeneetText.
OPT-125M
Method
PPL. 4
Mem.
Lat. |
Vanilla
19.82
408
4.88
Performer
61.20 (+41.38)
1.21
SEA (Ours)
22.64 (+2.82)
187
6.76
Validation Curve
150
125
Performer
100
PPL
75
2000 4000
6000
8000
10000
Optimizer Steps
Figure A.5: Validation curves for SEA and Performer, when trained with OpenWebtext dataset. SEA converges
much faster compared to Performer.
In Table A.7, we evaluate our method with the larger Open WebText (Gokaslan al., 2019) dataset. With a given training budget our method still outperforms the baseline, Performer. In this large- scale dataset, Performer shows much poorer approximation performance than our method. We train the models only 10k optimizer step. which is equivalent to 640M tokens. Considering the fact that LLMS are often trained with over 1T tokens (Zhang et al., 2023), we preserve most of the teacher performance with only 0.064% training cost of the teacher's pretraining, while also reducing computational complexity from quadratic to linear. Moreover, as shown Fig. A.5, SEA converges much faster compared to Performer, and the difference is even larger than Fig. 5b, where Wikitext2 dataset was used for training.
A.10 DETAILED DIAGRAM OF MODEL STRUCTURE
(1) Attention Estimationn
(2) Sparse Attention Mask Generation and Subsequent Sparse Attention
topa
self Amt
mai (A. 2) add par
F&
Selection Groupa
M
Decader
M
mix
San
Canol Amn
Spans
Weatordnn
M
A
Figure A.6: Model diagram of our proposed linear attention method, SEA, during inference.
In Fig. 2, we provide high-level overview of our method structure. And in this section, we provide a much more detailed model diagram with more notations in Fig. A.6. In the figure, (1) SEA estimates
21

---

Published as a conference paper at ICLR 2024
the attention matrix in compressed size (Á E RT*X). (2a) We then perform top-k selection procedure from the compressed attention matrix resulting in a compressed attention mask, and (2b) interpolate the compressed attention mask to sparse formatted mask for the full attention matrix. (2c) Finally, we perform sparse attention.
A.11 DETAILED RESULT OF SEA ON MNLI
Table A.8: Comparison with MNLI dataset of GLUE henchmark (Wang al., 2019) among linear attention methods. Vanilla quadratic attention included for reference. SEA (Ours) trained with 32 and &
Metric
Vanilla
SEA (Ours)
Cnsformer
Reformer
Sinkhorn
Synthesizer
Performer
Accuracy
84.1
84.0
82.7
82.5
81.9
75.5
74.7
Memory (MB)
9.00
17.17
10.88
88.36
9.39
8.25
14.76
Lateney (us)
238
701
242
900
152
181
320
In Table A.8, we provide the detailed result of SEA on MNLI which is the subset of the GLUE benchmark.
A.12 DISCUSSION ABOUT THE LATENCY OF OUR FLATCSR IMPLEMENTATION
We think there is room for further optimization of FlatCSR in order further reduce latencies. The- oretically, dense operations cost 43.96 GMACs, and sparse FlatCSR operation costs 1.25 GMACs, as shown in Fig. A.4; this means our kernel implementation does not fully utilize MACs and is therefore bottle-necked on memory computation and thread scheduling. Furthermore, we need to utilize block aligns in the sparse interpolation of the mask from M to M Since we are using the nearest neighborhood interpolation and we mostly interpolate the attention estimation from relatively smaller size, K, to larger size T, single non-zero pixel in the compressed mask will be multiple pixels in resized one. We may utilize this fact to schedule threads with blocks like other famous block-sparse attention implementations including FlashAttention Dao et al. (2022). There- fore, further research could look to investigating and optimizing thread scheduling and cache hit ratios of the proposed FlatCSR. However, think the implementation with Triton is sufficient to show the efficiency of the proposed SEA attention.
22


Reviews:
[{"review_id": "Z0NRiw2SvF", "reviewer": "Reviewer_qYba", "summary": "This work introduces SEA, a novel method to approximate the full dot-product attention at inference time. It uses two mechanisms: (i) it relies on a linearization to build a compressed attention representation and then (ii) generates a full-scale sparse attention mask from this representation. They show their method outperforms other linear and sparse attention methods in language modeling tasks, specifically on GLUE and Wikitext2, while being competitive in terms of memory usage and latency. SEA introduces additional model parameters which are trained on top of a pretrained model using knowledge distillation. Unlike some of its competitors, the method remains interpretable.", "questions": "- How would SEA compare to FlashAttention?\n- How difficult is it to tune the weights given to the different loss terms? \n- Wikitext is a relatively narrow dataset, how would the approximation handle more diverse datasets such as openwebtext?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Enabling faster processing of long sequences is an important research direction, and the proposed method is well-motivated. I appreciate the effort made in presenting the method, which, despite its complexity, can still be understood. The idea of combining kernel-based linear attention and sparsification is novel. On GLUE tasks, experiments show how SEA approximates full attention better than other methods while remaining competitive in terms of memory footprint. Moreover, unlike other approaches, SEA can successfully approximate the full-attention of pretrained OPT models.", "weaknesses": "- Comparison with FlashAttention [1]: It would be fair to add FlashAttention among the baselines. Especially, FlashAttention would also be competitive in terms of memory.\n- The method is still quite complex, making it hard to deploy. \n- The latency results do not show a clear advantage of the method over baselines, often being significantly slower.\n- The justification of the method for autoregressive language modeling is unclear. As most causal models are used for sequence generation, sampling one token at a time, the need to sparsify the attention matrix at inference time is reduced.\n\nReferences:\n[1] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness", "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": "", "extra_scores": {}}, {"review_id": "LLL62Cza3O", "reviewer": "Reviewer_VaUL", "summary": "The paper proposes a method to improve the quadratic dependency of the attention mechanism on the token length to linear, addressing a critical issue in computational complexity.\n\n\n\n\n\n------------------------------------------------------------------------------------\npost rebuttal update:\nI'd like to thank the authors for answering my questions. I raise my score.", "questions": "- Can your algorithm be applied to pre-trained models for faster inference without requiring fine-tuning? Can one apply your algorithm for faster inference with no fine-tuning whatsoever?\n\n- Could you provide a detailed explanation of Figure 2? I'm having difficulty understanding this diagram. Is the decoder (MLP CNN) trained from scratch, or can it be extracted from pre-trained models?\n\n- In general, do you train all components depicted in Figure 2, or do you incorporate some parts from pre-trained models while keeping them fixed?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "The computational complexity of the attention mechanism is a serious bottleneck and improving this to linear is very useful. The strength of the paper is that it tackles an important problem.", "weaknesses": "The paper's clarity and explanation of the algorithm's functionality are lacking, making it challenging to determine its applicability, particularly with regard to pre-trained models.", "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": "", "extra_scores": {}}, {"review_id": "7xJN1sQaLw", "reviewer": "Reviewer_V5oR", "summary": "The authors propose a new test-time sparse linear attention method that 1) estimates the quadratic attention in compressed size, in the style of Performer; 2) performs a top-k selection on the approximated attention to obtain a mask, and 3) interpolates top-k selection to obtain a full sparse attention matrix and then perform sparse attention. They show that their method achieves better perplexity on accuracies on WikiText2 and GLUE tasks respectively.\n\nThe authors also contribute a new Triton kernel for performing efficient sparse attention.", "questions": "Please, see W1-3 above. It would be great to revise your paper based on my comments and questions. Thanks.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "S1. I like the approach of a \"test-time\" sparse linear attention. \n\nS2. I appreciate the thorough description of the author's contributions.\n\nS3: I appreciate the contribution of a new Triton kernel for sparse operations.", "weaknesses": "W1. One of the motivations of the paper is that other linear attentions cannot distill the learned attention patterns, and hence need to train from scratch. However, the authors in the paper still need to train their Performer and Decoder from scratch. I haven't seen any discussion about the inherent cost of doing that. Intuitively, it should be cheaper than training from scratch, but can you point me to the text (or elaborate in a new discussion) about how expensive it is to do this training?\n\nW2. This is my subjective view, but the paper is extremely dense and hard to follow. I'd recommend reducing the notation significantly and moving most of the details to the appendix. I appreciate the figures, but overwhelming them with notation does aid my understanding of your method in the current version of the paper.\n\nI recommend: 1) simpler figures with less notation that can clearly explain your method conceptually; 2) moving a lot of the notation in the text for the appendix.\n\nW3. It seems to me that some of the results are a bit underwhelming. For example, in Figure 4, panel (a), right figure, what is the motivation to use your method, since I could use a Vanilla transformer and achieve an on-par accuracy with much less latency?", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)