#  SPION: Layer-Wise Sparse Training of Transformer via Convolutional Flood
Filling

Bokyeong Yoon  Sogang University35 Baekbeom-ro, Mapo-guSeoulRepublic of Korea
<EMAIL_ADDRESS>,  Yoonsang Han  Sogang University35 Baekbeom-ro, Mapo-
guSeoulRepublic of Korea<EMAIL_ADDRESS>and  Gordon Euhyun Moon  Sogang
University35 Baekbeom-ro, Mapo-guSeoulRepublic of Korea<EMAIL_ADDRESS>
###### Abstract.

Sparsifying the Transformer has garnered considerable interest, as training
the Transformer is very computationally demanding. Prior efforts to sparsify
the Transformer have either used a fixed pattern or data-driven approach to
reduce the number of operations involving the computation of multi-head
attention, which is the main bottleneck of the Transformer. However, existing
methods suffer from inevitable problems, such as the potential loss of
essential sequence features due to the uniform fixed pattern applied across
all layers, and an increase in the model size resulting from the use of
additional parameters to learn sparsity patterns in attention operations. In
this paper, we propose a novel sparsification scheme for the Transformer that
integrates convolution filters and the flood filling method to efficiently
capture the layer-wise sparse pattern in attention operations. Our
sparsification approach reduces the computational complexity and memory
footprint of the Transformer during training. Efficient implementations of the
layer-wise sparsified attention algorithm on GPUs are developed, demonstrating
a new SPION that achieves up to 3.08$\times$ speedup over existing state-of-
the-art sparse Transformer models, with better evaluation quality.

##  1\. Introduction

The Transformer is a state-of-the-art deep neural network developed for
addressing sequence tasks, originally proposed by Vaswani et al. (Vaswani et
al., 2017). One of the main advantages of the Transformer is that, given a
sequence of input data points (e.g., a sentence of word tokens), it is able to
compute the multi-head attention (MHA) operation in parallel, thereby quickly
and accurately capturing long-term dependencies of data points. However, as
the sequence length increases, it is true that the overall computational cost
and memory space required for training the Transformer also increase
quadratically (Ainslie et al., 2020). Especially, a MHA sub-layer in the
Transformer occupies a substantial portion of the total execution time and
becomes the main bottleneck as the sequence length increases. The MHA
operation requires a large number of dot-product operations to compute the
similarity between all data points in the sequence. However, the dot-product
operation inherently has limitations in memory bandwidth since it performs
only two floating-point operations for each pair of data elements read from
memory. Therefore, reducing the number of operations involving dot-product
operations is an important consideration for accelerating the training of the
Transformer.

In addition, the attention matrices used in the MHA computation have two main
challenges. First, as the pattern of non-zero entries in the attention
matrices varies across different layers, the irregular access of these non-
zero entries within the matrices leads to an inefficient utilization of global
memory bandwidth and low cache hit rates while performing dot-product
operations. Second, the irregular distribution of non-zero entries in
attention matrices across different layers results in load imbalance and a
reduced level of exploitable parallelism. Hence, in order to mitigate
computational complexity and improve data locality of the Transformer, several
approaches have addressed the sparsification of computations associated with
the MHA operation (Zaheer et al., 2020; Beltagy et al., 2020; Roy et al.,
2021; Zhang et al., 2021; Wang et al., 2020a). However, previous approaches
suffer from two primary limitations. Firstly, when the Transformer adopts
identical fixed patterns of non-zero entries in the attention matrices across
all layers during training, it becomes difficult to effectively capture key
features within the sequence. Secondly, when the Transformer employs
additional parameters to learn the sparsity pattern in the attention matrices,
both the model size and the computational overhead increase. To address the
limitations of previous approaches, we focus on a specialized sparse attention
algorithm that not only significantly reduces memory consumption but also
efficiently captures sparse patterns across various types of Transformers and
datasets.

In this paper, we present a new sparsity-aware layer-wise Transformer (called
SPION) that dynamically captures the sparse pattern in the MHA operation for
each layer. SPION judiciously explores the sparsity pattern in the MHA
operation based on a novel convolutional flood filling method. To precisely
detect the characteristics of the sparse pattern in the attention matrices,
SPION captures the shape of sparse pattern by utilizing a convolution filter
and the degree of sparsity through a flood filling-based scheme. During the
generation of the sparse pattern for each layer, our SPION constructs a
sparsity pattern matrix with a blocked structure, enhancing data locality for
the MHA operation that involves sparse matrix multiplication. Furthermore, by
capturing layer-specific sparse pattern for each layer, SPION performs layer-
wise MHA computations iteratively until convergence is achieved. Our
sparsification technique reduces both the time and the number of operations
required to train long sequences without using additional parameters.
Moreover, since the sparse MHA operations contribute significantly to the
overall training workload, we develop an efficient GPU implementation for
sparse MHA to achieve high performance with quality of results.

We conduct an extensive comparative evaluation of SPION across various
classification tasks, including image and text classification, as well as
document retrieval involving lengthy sequences. Experimental results
demonstrate that our SPION obtains up to 10$\times$ lower operations and up to
3.08$\times$ training speedup compared to existing state-of-the-art sparse
Transformers.

The paper is organized as follows. Section 2 presents the background on
Transformer, sparse attention, flood fill algorithm, and related prior work.
Section 3 describes the motivation behind our SPION. Section 4 presents an
overview and details of our SPION implementation. Section 5 extensively
compares the performance of SPION with various state-of-the-art sparse
Transformer models.

##  2\. Background and Related Work

###  2.1. Transformer

The encoder-only Transformer, which is one of the variants of the Transformer
model, is widely used for various classification tasks using text and image
datasets (Tay et al., 2022; Devlin et al., 2019; Dosovitskiy et al., 2021).
Algorithm 1 shows pseudo-code for the forward propagation of the original
encoder-only Transformer, where each encoder layer consists of a MHA sub-layer
(lines 2-10) and a feed-forward sub-layer (lines 11-12).

Input: $E\in\mathbb{R}^{L\times D}$: a set of embedding vectors for an input
sequence, $L$: length of input sequence , $D$: embedding size for each data
point, $H$: the number of heads, $W^{Q}$, $W^{K}$, $W^{V}$, $W^{O}$, $W^{F}$,
$W^{E}$: weight parameters, $N$: the number of encoder layers

1

2for  _$encoder\\_layer=0$ to $N-1$_ do

$\triangleright$ MHA sub-layer

3 $X\leftarrow LayerNorm(E)$

4 $Q\leftarrow X\times W^{Q},K\leftarrow X\times W^{K},V\leftarrow X\times
W^{V}$

5 $Q_{0,...,H-1}$, $K_{0,...,H-1}$, $V_{0,...,H-1}$ $\leftarrow split(Q,K,V)$

6 parallel for  _$i=0$ to $H-1$_

7

8 $A_{i}^{r}\leftarrow Q_{i}\times K_{i}^{T}$

9 $A_{i}^{s}\leftarrow softmax(A_{i}^{r}/\sqrt{D/H})$

10 $A_{i}^{c}\leftarrow A_{i}^{s}\times V_{i}$

11 $A\leftarrow concatenate(A^{c}_{0},...,A^{c}_{H-1})$

12 $O\leftarrow dropout(A\times W^{O})+E$

$\triangleright$ Feed-forward sub-layer

13 $F\leftarrow ReLU(LayerNorm(O)\times W^{F})$

14 $E\leftarrow dropout(F\times W^{E})+O$

Algorithm 1 Forward propagation of encoder layer in encoder-only Transformer

For each encoder layer, the input embedding $E$ is transformed into $X$ by
applying layer normalization (line 2), and the query ($Q$), key ($K$), and
value ($V$) are obtained by performing linear transformations on the embedding
$X$ (line 3). Hereafter, we denote $L$, $D$ and $H$ as the length of input
sequence, the size of the embedding for each data point in the sequence, and
the number of heads, respectively. To efficiently perform MHA computation,
each $Q\in\mathbb{R}^{L\times D}$, $K\in\mathbb{R}^{L\times D}$, and
$V\in\mathbb{R}^{L\times D}$ matrix is divided into $H$ sub-matrices (multi-
heads) along the $D$ dimension (line 4). The utilization of multi-heads
indicates that the embedding of the input sequence is divided into multiple
subspaces. This division enables the model to perform MHA in parallel (lines
5-8) while precisely capturing distinct key features in each subspace. The
loop in lines 5-8 performs the MHA computation as formulated in Equation 1.

(1) $A=softmax\left(\frac{Q\times K^{T}}{\sqrt{(D/H)}}\right)\times V$

The computation within $i$ loops involves matrix-matrix multiplication (GEMM)
to obtain the raw attention score matrix $A_{i}^{r}$, which captures
similarity between query $Q_{i}\in\mathbb{R}^{L\times(D/H)}$ and key
$K_{i}\in\mathbb{R}^{L\times(D/H)}$ for head $i$ (line 6). The softmax
function is then applied to the matrix $A_{i}^{r}$, which is scaled by
$1/\sqrt{D/H}$ (line 7). This scaling is performed to prevent the gradients of
the attention score matrix $A_{i}^{s}$ from approaching zero. Thereafter,
another GEMM operation is performed between matrices $A_{i}^{s}$ and
$V_{i}\in\mathbb{R}^{L\times(D/H)}$ to obtain the complete attention matrix
$A_{i}^{c}$ for head $i$ (line 8). After computing the attention for each
head, all matrices $A_{0}^{c}$ through $A_{H-1}^{c}$ are concatenated to form
the final attention matrix $A$, and the input embedding $E$ used at the
current encoder layer is added to the attention matrix (lines 9 and 10). Then,
the attention matrix is passed through the feed-forward sub-layer to produce
new embedding vectors $E$, which is then fed into the next encoder layer
(lines 11 and 12). In terms of computational complexity, it is clear that the
main bottleneck of the encoder layer is associated with processing the MHA
sub-layer (lines 2-10). More specifically, the number of operations required
for computing the attention matrix $A$ for head $i$ (loop in line 5) is
$2L^{2}(2D+1)-L(D+1)$, indicating a quadratic increase in operations as the
input sequence length ($L$) increases.

Sparsification of MHA. Given the long input sequences, several sparse
attention techniques have been proposed to reduce the computational complexity
involved in computing the raw attention matrix $A^{r}$. The basic intuition
behind sparse attention techniques is that a subset of data points in the long
sequence can effectively represent the entire input sequence. In other words,
to reduce the computational workload of computing the raw attention matrix
$A^{r}$, only the highly correlated necessary elements (i.e., data points) in
$Q$ and $K$ can be utilized. Therefore, when employing sparse attention, the
number of operations required to compute $A_{i}^{r}$ for each head $i$ is
$C\times(2D-1)$, where $C$ is the number of critical data points in the long
sequence, whereas the number of operations required for the original dense
$A_{i}^{r}$ is $L^{2}\times(2D-1)$.

###  2.2. Flood Fill Algorithm

The flood fill scheme is originally developed to determine the area connected
to a given cell/pixel in a multi-dimensional array and a bitmap image
(Goldman, 1990). Starting from a given seed element, the algorithm determines
whether to continue or stop the fill operation based on conditional statements
that evaluate the properties of the current element and its surrounding
element. If the current element is not at the end of the data structure and
the neighboring elements have not been explored yet, the algorithm continues;
otherwise, it is avoided. Furthermore, the flood fill algorithm dynamically
controls the fill operation with a constraint on small values of the elements.
If the value of current element is very small, it prevents the fill operation.

###  2.3. Related Work on Sparse Attention

Many previous efforts to achieve efficient Transformers have sought to
sparsify the MHA operation both before and during model training/fine-tuning.

Fixed Sparse Pattern. One of the strategies for performing sparse MHA is to
use a predetermined sparsity patterns, where only specific data points in the
input sequence are selected before training the model. In this approach, only
these selected data points are exclusively used to perform the MHA operation.
Several variants of the encoder-only Transformer model adapt the sliding
windows approach in which the attention operations are performed using only
the neighboring data points (rows) in the matrices $Q$ and $K$. The Sparse
Transformer (Child et al., 2019) originally employs the sliding windows
attention to sparsify the MHA operation. The Longformer (Beltagy et al., 2020)
is an extension to the Sparse Transformer and introduces dilated sliding
windows, which extend the receptive field for computing similarity by skipping
one data point at a time while performing the sliding windows attention.
Furthermore, ETC (Ainslie et al., 2020) and BigBird (Zaheer et al., 2020)
incorporate global attention that performs similarity calculations between a
given data point and all other data points in the input sequence. ETC
introduces three types of fixed sparsity patterns for attention, called
global-global, global-local, and local-global. BigBird further proposes random
attention, which randomly selects data points for similarity computation.

The main advantage of sparsifying the MHA operation before training is the
reduction in computational overhead. As the fixed sparsity patterns are
applied to all MHA computations in the entire training procedure, this
approach enables reducing the overall memory footprint compared to the
original dense MHA operation. However, the primary problem with a fixed
pattern is that it is not sufficient to fully capture the dependencies present
in the input sequence. For example, depending on the types of dataset and
task, the critical data points in the long input sequence can vary. Therefore,
predefined static sparse patterns may lose the fine-grained important features
and dependencies in the sequence during training.

Data-driven Sparse Pattern. Sparse patterns in the MHA operation can also be
generated by leveraging a data-driven approach, which clusters and sorts the
data points of the input sequence during training. For example, Reformer
(Kitaev et al., 2020) utilizes locality-sensitive hashing to calculate a hash-
based similarity and cluster multiple data points into chunks. Similarly,
Routing Transformer (Roy et al., 2021) performs k-means clustering given data
points. In the process of training the model, these clustering-based attention
approaches also require learning additional functions to effectively identify
the relevant dependencies of data points in the input sequence. Wang et al.
(Wang et al., 2020b) reduce the number of operations required for MHA
computation by parameterizing weight matrices involved in MHA using low-rank
matrix factorization and eliminating rank-1 components during model training.
Peer et al. (Peer et al., 2022) and Lagunas et al. (Lagunas et al., 2021)
reduce the size of the fine-tuning model by removing the insignificant layers
and sub-blocks in attention matrix based on the parameters from the pretrained
model. However, even though utilizing data-driven sparsification techniques
during training produces a better quality model, this approach requires
additional parameters and operations to learn the sparse patterns, resulting
in larger memory space and higher computational cost.

##  3\. Motivation: Analysis of Sparse Patterns in MHA

Figure 1. Sparsity patterns in the attention score matrices $A^{s}$ across
different encoder layers during training of the encoder-only Transformer for
image classification.

(a) Overall training/fine-tuning procedure of SPION

(b) Dense MHA operation

(c) Sparse MHA operation

Figure 2. SPION splits the overall training process into three phases: dense-
attention training with dense MHA operation, sparsity pattern generation, and
sparse-attention training with sparse MHA operation

In order to recognize the common sparsity patterns in MHA operation, we
conducted experiments on the pretrained encoder-only Transformer (Dosovitskiy
et al., 2021). Note that the pretrained model is trained on a large
ImageNet-21K dataset. Figure 1 illustrates the attention score matrices
$A^{s}$ from different encoder layers during the first epoch of fine-tuning
the model with the CIFAR-10 dataset. Since the sparsity patterns of multiple
attention score matrices within the same encoder layer typically show similar
patterns, we averaged the attention score matrices across multiple heads in
each encoder layer. The results clearly show that most of the elements in
$A^{s}$ are close to zero, indicating that only a few data points in the input
sequence are correlated to each other and considered critical. In practice, as
the sparse pattern of matrix $A^{s}$ is maintained for each layer after
training the model for a few epochs, a number of studies have shown that
considering only the critical data points does not adversely affect the
convergence of the model (Beltagy et al., 2020; Zaheer et al., 2020; Qiu et
al., 2020). The characteristics of attention score matrices $A^{s}$ observed
from our experiments are as follows.

Shape of Sparse Pattern. As shown in Figure 1, the attention score matrices
$A^{s}$ produced by the different encoder layers exhibit distinct sparsity
patterns. For example, in the first to eighth encoder layers, the diagonal
elements have relatively large values, similar to a band matrix that stores
nonzeros within the diagonal band of the matrix. It is obvious that the MHA
operation relies on the self-attention scheme and therefore, the resulting
values of the dot-product between linearly transformed vectors for the same
data points tend to be larger compared to the resulting outputs produced with
different data points. In addition to the diagonal sparsity pattern, some
encoder layers, specifically layers 9, 10, 11 and 12, show a vertical sparsity
pattern, with nonzeros mostly stored in specific columns. This vertical
sparsity pattern emerges when the attention operation focuses on the
similarity between all data points in $Q$ and particular data points in $V$.
In light of these observations, applying the same fixed sparse pattern to all
layers may lead to the exclusion of unique essential features that need to be
captured individually at different layers. Hence, during the training of the
Transformer model, to efficiently reduce the number of operations, it is
crucial to consider layer-wise sparsification of the MHA based on the sparse
pattern observed across different layers. Furthermore, the flexibility of
changing the sparse pattern needs to be considered during the training of the
model. Moreover, generating domain-specific sparse patterns is required for
various datasets and tasks.

Degree of Sparsity. Across different encoder layers, there exists variation
not only in the shape of the sparse pattern but also in the number of nonzero
elements in $A^{s}$. For example, layer 12 has a higher number of nonzero
elements compared to layer 2, indicating that layer 12 extensively computes
the MHA operation using a larger number of data points in the sequence. Hence,
it is crucial to consider varying degrees of sparsity for every encoder layer.
This approach is essential for effectively reducing computational operations
while preserving key features across distinct encoder layers.

##  4\. SPION: Layer-Wise Sparse Attention in Transformer

###  4.1. Overview of SPION

Input: $E$: input embedding vectors, $L$: length of input sequence, $D$:
embedding size, $N$: the number of encoder layers, $\alpha$: threshold for
transition

1

2$transition\leftarrow False$

3for  _$i=0$ to $num\\_epochs-1$_ do

4

$\triangleright$ Dense MHA Phase

5 if  _transition == False_ then

6

7 for  _$n=0$ to $N-1$_ do

8

9 $O,A_{i}^{s}\leftarrow MHA(E)$

10 $E\leftarrow Feed\\_Forward(O)$

11 if  _i ¿ 1_ then

12

13 $distance_{i-1}\leftarrow frobenius(A_{i-2}^{s},A_{i-1}^{s})$

14 $distance_{i}\leftarrow frobenius(A_{i-1}^{s},A_{i}^{s})$

15 if  _$\sqrt{(distance_{i-1}-distance_{i})^{2}} <\alpha$_ then

16

17 $transition\leftarrow True$

18 $P\leftarrow generate\\_pattern(A_{i}^{s},t)$

$\triangleright$ Sparse MHA Phase

19

20 if  _transition == True_ then

21

22 for  _$n=0$ to $N-1$_ do

23

24 $O\leftarrow sparseMHA(E,P_{n})$

25 $E\leftarrow Feed\\_Forward(O)$

Algorithm 2 Overall Training Procedure of SPION Figure 3. Overview of our
convolutional flood filling method in the sparsity pattern
generation/detection phase

In this section, we provide a high-level overview and details of our new SPION
that dynamically sparsifies the MHA operation, incorporating the major
considerations described in Section 3. Note that our SPION is capable of
sparsifying the MHA operation either after training the model for a few steps
or during the process of fine-tuning the model from the scratch using the pre-
trained network.

As shown in Figure 2(a), we decouple the overall training process of SPION
into three phases: dense-attention training, sparsity pattern generation, and
sparse-attention training. The dense-attention training follows the same
training procedure as the original Transformer, without sparsifying the MHA
operation. During the dense-attention training phase, the original MHA, which
includes two GEMM kernels and a dense softmax kernel (Figure 2(b)), is
performed until the attention score matrix $A^{s}$ of each encoder layer
exhibits a specific sparsity pattern. In order to determine the end of the
dense-attention phase or the start of the sparse-attention phase, we measure
the Frobenius distance between the attention score matrices $A^{s}$ produced
in the previous step $i-1$ and the current step $i$ as defined in Equation 2.

(2)
$distance_{i}=\bigg{|}\sqrt{\sum(A^{s}_{i-1})^{2}}-\sqrt{\sum(A^{s}_{i})^{2}}\bigg{|}$

If there is no significant difference in distance, our SPION ceases the dense-
attention training phase and dynamically generates the sparsity pattern for
each encoder layer based on our novel convolutional flood-filling scheme,
which is described in Section 4.2. Intuitively, it can be assumed that the
attention score matrix is generalized when the values in the current attention
score matrix are similar to those in the previous attention score matrix. This
indicates that it is ready to apply sparsification to the attention score
matrix. Thereafter, SPION proceeds with the sparse-attention training phase
until convergence by adapting the sparsity pattern, as depicted in Figure
2(a). Given the dense matrices $Q$ and $K$, along with the newly generated
sparsity pattern matrix $P$, the SDDMM (Sampled Dense-Dense Matrix
Multiplication) kernel is utilized to accelerate producing the sparsified
attention score matrix $S^{r}$. Next, the sparse matrix $S^{r}$ is used to
perform the optimized sparse softmax kernel, leveraging warp-level reduction
for accelerating performance. After applying the sparse softmax operation,
since the attention score matrix $S^{s}$ remains sparse, we utilize SpMM
(Sparse-Dense Matrix Multiplication) kernel to multiply the sparse matrix
$S^{s}$ with the dense matrix $V$ (Figure 2(c)).

###  4.2. Sparsity Pattern Generation with Convolutional Flood Fill Algorithm

Input: $A^{s}\in\mathbb{R}^{L\times L}$, $t$ : threshold value for flood fill
function

Output: $P$: sparsity pattern matrix

1

2$filter\leftarrow\text{generate diagonal filter: }(F\times F)$

3$conv\\_out\leftarrow\text{convolution}(A^{s},filter)$

4$pool\\_out\leftarrow\text{avgpool}(conv\\_output,(B,B))$

5$fl\\_out\leftarrow$ initialize with zeros : $(L/B\times L/B)$

6for  _$i=0$ to $L/B-1$_ do

7
$pool\\_out,r,c,fl\\_out,t\leftarrow\text{flood\\_fill}(pool\\_out,0,i,fl\\_out,t)$

8for  _$j=0$ to $L/B-1$_ do

9
$pool\\_out,r,c,fl\\_out,t\leftarrow\text{flood\\_fill}(pool\\_out,j,0,fl\\_out,t)$

10for  _$k=0$ to $L/B-1$_ do

11 $fl\\_out[k,k]\leftarrow 1$

$P\leftarrow\text{upsampling}(fl\\_out,(B,B))$

Algorithm 3 Implementation of generate_pattern() function

To precisely detect the sparsity patterns in the attention score matrix
$A^{s}$ generated during dense-attention training, we develop a new
convolutional flood fill algorithm that extensively explores the shape, degree
and locality of sparsity patterns for each encoder layer. Algorithm 3 shows
the pseudo-code for generating the sparsity pattern in the attention score
matrix $A^{s}$. An initial step is to apply a diagonal convolution filter to
$A^{s}$ (line 1 in Algorithm 3) in order to detect the shape of sparsity
pattern in $A^{s}$, as shown in Figure 3. If the diagonal elements in $A^{s}$
have larger values compared to the others, applying a diagonal convolution
filter increases the values of the diagonal elements in the convolution output
($conv\\_out$), leading to the emergence of a diagonal sparsity pattern.
Otherwise, if the off-diagonal elements, especially the vertical ones, in
$A^{s}$ have larger values compared to the others, applying a diagonal
convolution filter results in a vertical sparsity pattern in $conv\\_out$
matrix. Note that in order to ensure that the attention score matrix $A^{s}$
and the convolution output ($conv\\_out$) have the same size ($L\times L$), we
adopt zero-padding to $A^{s}$ during computing the convolution operation
defined in Equation 3.

(3) $conv\\_out(i,j)=\sum_{f=1}^{F}A^{s}(i+f,j+f)\times filter(f,f)$

After generating the $conv\\_out$ through a diagonal convolution operation,
our algorithm performs average pooling on the $conv\\_out$ using a
kernel/block of size ($B\times B$) as defined in Equation 4 (line 2 in
Algorithm 3).

(4)
$pool\\_out\left(\frac{i}{B},\frac{j}{B}\right)=\frac{1}{B^{2}}\sum_{p=1}^{B}\sum_{q=1}^{B}conv\\_out(i+p,j+q)$

Instead of analyzing the sparsity pattern of $A^{s}$ element by element,
applying average pooling enables capturing block sparsity pattern, which takes
into account both the critical data points and their surrounding data points.
Hence, since the output of average pooling ($pool\\_out$) has a smaller size
($L/B\times L/B$) compared to the attention score matrix $A^{s}$, $pool\\_out$
can be considered as an abstract sparsity representation of $A^{s}$ in blocks.

Input: $pool\\_out\in\mathbb{R}^{L/B\times L/B}$, $r$: current row index, $c$:
current column index, $fl\\_out\in\mathbb{R}^{L/B\times L/B}$, $t$: threshold
value

1

2if  _$(r+1==L/B)$ or $(c+1==L/B)$_ then

3 return $pool\\_out$, $r$, $c$, $fl\\_out$, $t$

4$m\leftarrow\text{max}(pool\\_out[r+1][c],pool\\_out[r][c+1],pool\\_out[r+1][c+1])$

5if  _$pool\\_out[r+1][c]==m$ and $fl\\_out[r+1][c]==0$_ then

6

7 if  _$pool\\_out[r+1][c] >t$_ then

8

9 $fl\\_out[r+1][c]\leftarrow 1$

10
$pool\\_out,r,c,fl\\_out,t\leftarrow\text{flood\\_fill}(pool\\_out,r+1,c,fl\\_out,t)$

11if  _$pool\\_out[r][c+1]==m$ and $fl\\_out[r][c+1]==0$_ then

12

13 if  _$pool\\_out[r][c+1] >t$_ then

14

15 $fl\\_out[r][c+1]\leftarrow 1$

16
$pool\\_out,r,c,fl\\_out,t\leftarrow\text{flood\\_fill}(pool\\_out,r,c+1,fl\\_out,t)$

17if  _$pool\\_out[r+1][c+1]==m$ and $fl\\_out[r+1][c+1]==0$_ then

18

19 if  _$pool\\_out[r+1][c+1] >t$_ then

20

21 $fl\\_out[r+1][c+1]\leftarrow 1$

22
$pool\\_out,r,c,fl\\_out,t\leftarrow\text{flood\\_fill}(pool\\_out,r+1,c+1,fl\\_out,t)$

return $pool\\_out$, $r$, $c$, $fl\\_out$, $t$

Algorithm 4 Implementation of flood_fill() function

In order to dynamically explore the crucial elements in the $pool\\_out$ and
generate $fl\\_out$, we develop a novel algorithm inspired by the flood-fill
algorithm. By adapting the flood filling scheme, our SPION is able to
precisely analyze the connectivity between significant elements in
$pool\\_out$, while also considering the number of critical nonzero elements.
Algorithm 4 shows the pseudo-code for our flood fill-based algorithm
recursively executed in Algorithm 3. Unlike the traditional flood fill
algorithm, which compares all neighbors of a current element to find the
element with the largest value, our new algorithm only compares to the
neighboring elements on the right, below, and diagonally below, as shown in
Figure 4. In the process of capturing the pattern, it is necessary to
sequentially follow the important features starting from the top-left corner
to the bottom of the matrix. Hence, comparing the current element with the
elements to its right or below checks whether the neighbors of the current
element are relevant. If the neighboring elements are not relevant to the
current element, our algorithm moves to the element diagonally below to
continue comparing the neighbors.

Figure 4. Walk-through example of our flood fill algorithm

The element at the first row and first column of the $pool\\_out$ is used as
the starting point (lines 4 and 6 in Algorithm 3). Then the algorithm compares
the values of elements to its right, below, and diagonally below, extracting
the element with the largest value, $m$ (line 3 in Algorithm 4). When the
value of $m$ is greater than a predefined threshold $t$ (line 5, 9, and 13 in
Algorithm 4), we determine the corresponding element as the critical element
and assign a value of 1 to the corresponding element in $fl\\_out$ (lines 6,
10 and 14 in Algorithm 4). Here, the threshold $t$, is determined by
calculating the $\alpha\%$ quantile of $pool\\_out$. Hence, even if one of the
neighbors is selected as a potential critical point, it may be classified as a
non-critical point based on the threshold $t$. By utilizing the threshold $t$,
we ensure that the selected critical elements have values that are
sufficiently large. After determining the critical element, the algorithm
recursively compares the values of elements to the right, below, and
diagonally below the critical element while avoiding duplicate comparisons
with elements that have already been selected. (second condition of lines 4,
8, and 12). This process is repeated until the selected critical element
reaches the end of a row or column in the matrix $pool\\_out$ (line 1 in
Algorithm 4). Afterward, as shown in the middle of Figure 4, the element at
the first row and second column serves as the next starting point for
recursively analyzing the connectivity of critical elements. Therefore,
considering each element in $pool\\_out$ as a seed point allows for the
exploration of the connectivity of critical elements across all elements in
$pool\\_out$. Since the attention score between the same data points in the
sequence tends to be large for most of the encoder layers, we initially assign
a value of 1 to the diagonal elements in $fl\\_out$.

The final output of our flood fill algorithm ($fl\\_out$) can be seen as a
more explicit sparsity pattern captured from $pool\\_out$ and can also be
considered as the compressed block-level sparsity pattern of $A^{s}$. To
utilize the sparse pattern of $fl\\_out$ in the sparse MHA operation, the size
of the matrix $fl\\_out$ needs to be the same as the size of the attention
score matrix $A^{s}$. Therefore, we upsample the $fl\\_out$ using nearest-
neighbor interpolation (line 11 in Algorithm 3), resulting in each nonzero
element in $fl\\_out$ forming a block of nonzero elements in the final
sparsity pattern matrix ($P$), as shown on the right of Figure 3. Utilizing a
block sparse matrix can improve the model’s quality since it incorporates not
only the critical elements but also their surrounding elements for MHA
operation. Moreover, an optimized blocked matrix multiplication further
enhances performance through improved data locality. In the blocked sparse
matrix $P$, elements that require calculation in the attention score matrix
are set to 1, while those that do not require calculation are set to 0.
Finally, during the sparse MHA, only the elements of $A^{s}$ that have the
same indices as the elements with a value of 1 in the block sparsity pattern
matrix $P$ will be computed.

###  4.3. Acceleration of Sparse MHA Implementation on GPUs

In this sub-section, we provide details of our parallel sparse MHA
implementation on GPUs. Prior to the sparse-attention phase and sparsity
pattern generation, we train the model with the dense MHA for several steps.
During the dense-attention phase, our SPION implementation uses the NVIDIA
cuBLAS library with the tensor cores on GPUs, such as
cublasGemmStridedBatchedEx(), to accelerate the dense MHA operation, which
involves the multiplication of dense matrices $Q$ and $K$, as well as the
multiplication of dense matrices $A^{s}$ and $V$.

Input: $E$: a set of embedding vectors for an input sequence, $L$: length of
input sequence , $D$: embedding size for each data point, $H$: the number of
heads, $W^{Q}$, $W^{K}$, $W^{V}$, $W^{O}$: weight parameters, $P$: sparsity
pattern matrix

1

2$X\leftarrow LayerNorm(E)$

3$Q\leftarrow X\times W^{Q},K\leftarrow X\times W^{K},V\leftarrow X\times
W^{V}$

4$Q_{0,...,H-1}$, $K_{0,...,H-1}$, $V_{0,...,H-1}$ $\leftarrow split(Q,K,V)$

5parallel for  _$i=0$ to $H-1$_

6

7 $S_{i}^{r}\leftarrow cusparseSDDMM(Q_{i},K_{i}^{T},P)$

8 $S_{i}^{s}\leftarrow SparseSoftmax(S_{i}^{r},P)$

9 $S_{i}^{c}\leftarrow cusparseSpMM(S_{i}^{s},V_{i},P)$

10$S\leftarrow concatenate(S^{c}_{0},...,S^{c}_{H-1})$

11$O\leftarrow dropout(S\times W^{O})+E$

Algorithm 5 GPU Implementation of sparseMHA() on host

Sparse MHA. Algorithm 5 shows the pseudo-code for the forward propagation of
sparse MHA. To accelerate the sparse MHA operation, we utilize NVIDIA cuSPARSE
libraries such as cuspraseSDDMM and cusparseSpMM. Additionally, we optimize
the softmax function to account for sparsity in the sparse raw attention score
matrix.

(5) $S=softmax\left(\frac{(P>0)\odot(Q\times
K^{T})}{\sqrt{(D/H)}}\right)\times V$

As shown in Equation 5, the sparse MHA produces a sparse attention matrix $S$
by replacing the GEMM operation in line 7 of Algorithm 1 with the SDDMM
operation using cusparseSDDMM (line 5 in Algorithm 5). The SDDMM kernel
computes the product of two dense input matrices $Q$ and $K$, and the
resulting matrix is then subjected to a Hadamard product (element-wise
multiplication) using a sampled sparse matrix $P$. In Equation 5, $(P>0)$
indicates that only the indices of nonzero elements in $P$ are computed in the
result matrix while multiplying $Q$ and $K$, and $\odot$ denotes element-wise
multiplication. Therefore, the SDDMM operation produces the sparsified raw
attention matrix $S^{r}=(P>0)\odot(Q\times K^{T})$, which stores only a small
number of nonzero elements. To efficiently compute $S^{s}$ by performing the
softmax operation with the sparse matrix $S^{r}$, we implement the sparse
softmax function (line 6 in Algorithm 5). This function computes the
probability distribution for only the nonzero values in $S^{r}$, instead of
applying the standard softmax function used in line 8 of Algorithm 1.
Thereafter, as $S^{s}$ remains a sparse matrix, we perform the SpMM operation
using cusparseSpMM to compute the product of a sparse matrix $S^{s}$ and a
dense matrix $V$ (line 7 in Algorithm 5). To perform cursparseSDDMM,
SparseSoftmax, and SpMM kernels with $P$, we convert the sparse matrix $P$
into the most commonly used Compressed Sparse Row (CSR) format consisting of
three data structures: $row\\_ptr$, $col\\_idx$ and $values$.

Input: $S^{r},P,scale$

1

2$warp\\_id\leftarrow\lfloor threadIdx.x/warp\\_size\rfloor$

3$lane\leftarrow threadIdx.x\%warp\\_size$

4$b\\_cnt\leftarrow P.row\\_ptr[(warp\\_id+1)\%L]-P.row\\_ptr[warp\\_id\%L]$

5$b\\_idx\leftarrow\lfloor warp\\_id/L\rfloor\times
P.values+P.row\\_ptr[warp\\_id\%L]$

6$max\leftarrow-\infty$

7$sum\leftarrow 0$

8for  _$k$ = $lane$ to $b\\_cnt$-1 by $warp\\_size$_ do

9 $S^{r}[b\\_idx+k]\leftarrow S^{r}[b\\_idx+k]\times scale$

10 if  _$max <S^{r}[b\\_idx+k]$_ then

11 $max\leftarrow S^{r}[b\\_idx+k]$

12$max\leftarrow warp\\_reduce\\_max(max)$

13for  _$k$ = $lane$ to $b\\_cnt$-1 by $warp\\_size$_ do

14 $sum\leftarrow sum+exp(S^{r}[b\\_idx+k]-max)$

15$sum\leftarrow warp\\_reduce\\_sum(sum)$

16$sum\leftarrow sum+exp(-max)\times(L-b\\_cnt)$

17for  _$k$ = $lane$ to $b\\_cnt$-1 by $warp\\_size$_ do

18 $S^{r}[b\\_idx+k]\leftarrow exp(S^{r}[b\\_idx+k]-max)/sum$

Algorithm 6 SparseSoftmax() kernel for forward propagation on GPUs

Sparse Softmax Kernel. Algorithm 6 shows the pseudo-code for our
SparseSoftmax() kernel. In the SparseSoftmax() kernel, as the original softmax
function operates in a row-wise manner, each warp is responsible for computing
a single row. The variable $warp\\_id$ indicates which row should be
calculated on the current warp, and the variable $lane$ indicates the lane
number within a single warp. For each row, $b\\_cnt$ indicates the number of
elements that should compute, and $b\\_idx$ indicates where the each row
starts at $S^{r}$. We initialize the variables $sum$ and $max$ to 0 and
$-\infty$, respectively. Since there is no need to share these variables among
different threads, the variables $sum$ and $max$ are initialized inside the
kernel. To ensure numerical stability, we normalize the elements in each row
by subtracting the maximum value found in each row. To find the maximum value
in each row, a lane in a warp fetches and compares a subset of elements within
that row (lines 9-11 in Algorithm 6). The function $warp\\_reduce\\_max()$ in
line 11 enables a lane to exchange values using warp-level primitives and
compare them with each other in a reduction manner. After executing the
$warp\\_reduce\\_max()$ function, each lane within a single warp receives the
same maximum value that was passed by all the lanes. We also use the warp-
level primitive to aggregate sum value across all lanes within a single warp
(line 14) and obtain the final value of $sum$. Finally, the computation of
sparse softmax function is completed by dividing the elements in the row by
the value of $sum$ variable. Therefore, by only considering the nonzero
elements in $S^{r}$, our SparseSoftmax() kernel greatly reduces the
computational complexity of softmax operation.

Integration of GPU kernels with PyTorch. To execute our optimized CUDA kernels
and CUDA libraries for both forward and backward propagation within the
PyTorch code, we integrate these kernels and libraries with PyTorch by
employing the ctypes library (CDLL) in the Python environment. Furthermore,
our custom autograd function is used to process the backward propagation.
Before the training begins, we first compile our CUDA code into a shared
object (.so) file. Once compiled, the CUDA shared object file is loaded and
then converted into a Python handle. Subsequently, the Python handle is
employed to invoke the CUDA kernels for both forward and backward
propagations.

###  4.4. Computational Complexity Analysis Comparison

By applying our sparsification scheme, the sparse raw attention matrix
$S_{i}^{r}$ maintains a total of $C$ non-zero elements, whereas the original
dense matrix $A_{i}^{r}$ maintains a total of $L\times L$ non-zero elements.
By using the sparse matrix $S_{i}^{r}$, the total number of operations
required to compute $S$ is reduced to $2C(2D+1)-L(D+1)$, resulting in
approximately $\frac{L^{2}}{C}\times$ reduction in operations compared to the
original attention matrix $A$, which requires $2L^{2}(2D+1)-L(D+1)$
operations. In practice, for the real-world AAN dataset (Radev et al., 2013)
($L=4,096$ and $D=64$) with the number of $ncd=1,677,721$ (10% of $L^{2}$),
the total number of operations required for producing the original $A$ and the
sparsified $S$ are $4,328,255,488$ and $432,585,778$, respectively. Therefore,
sparse MHA requires approximately 10 times less operations than the original
dense MHA during training.

##  5\. Experimental Evaluation

This section provides both performance and quality assessments of our SPION
implementation. Our SPION is compared with various state-of-the-art sparse
Transformer models.

Table 1. Machine configuration

Machine | Details
---|---
CPU |  | AMD Ryzen Threadripper PRO 5955WX
---
(16 cores and 32 threads, 128GB RAM)
GPU |  | NVIDIA RTX A5000
---
(24 GB Global Memory, 64 SMs, 6 MB L2 cache)

Benchmarking Machines. The detailed configuration of the benchmarking machines
used for the evaluations is shown in Table 1. All the experiments were run on
four NVIDIA RTX A5000 GPUs.

Datasets and Tasks. We evaluated our model on the Long Range Arena (LRA) (Tay
et al., 2020), which is widely used for evaluating the effectiveness of
efficient transformers for long sequences. We conducted three sequence
classification tasks as part of the LRA evaluations, including image
classification, ListOps and document retrieval tasks. For all experiments, we
measured classification accuracy to compare the quality of the compared
models. The three tasks, along with the three datasets used in our
evaluations, are as follows:

  * •

Image classification: This task classifies images into 10 different classes
using the CIFAR-10 dataset (Krizhevsky et al., 2009). Each image in the
dataset has a size of 32x32 pixels, resulting in a sequence length of 1,024 by
considering each pixel as a data point.

  * •

ListOps: This task involves classifying the answer within the range of 0 to 9,
given input equations expressed as a sequence of numbers and mathematical
operators. The dataset provided from Nangia et al. (Nangia and Bowman, 2018)
is used for evaluating this task. The maximum sequence length is 2,048.

  * •

Document retrieval: This task classifies whether two given documents are
related or not. The AAN dataset (Radev et al., 2013) is used for this task,
and the maximum sequence length is is 4,096.

Models Compared. We compared our SPION with the original encoder-only
Transformer and two state-of-the-art efficient sparse Transformers.

  * •

Original encoder-only Transformer (Vaswani et al., 2017): This implementation
is based on the original Transformer architecture and performs the original
dense MHA operation during entire training.

  * •

BigBird (Zaheer et al., 2020): This model incorporates sparse attention
mechanisms, including sliding window attention, global attention, and random
attention. It is evaluated using a block size of 64 and 3 random blocks.

  * •

Reformer (Kitaev et al., 2020): This model utilizes a sparse attention
mechanism based on locality-sensitive hashing algorithm. It is evaluated using
a bucket size of 32 and 2 hashes.

  * •

SPION-C: This model is the variation of our SPION model without applying the
flood filling scheme. Instead, it selects the top $\alpha\%$ of block elements
from the results of the convolution filter and average pooling ($pool\\_out$),
which enables adjustment of the sparsity ratio when generating the $P$.

  * •

SPION-F: This model is the variant of our SPION without using a convolution
filter. The model directly applies the flood-fill algorithm immediately after
the average pooling operation without applying a convolution filter.

  * •

SPION-CF: This implementation incorporates both the convolution filter and the
flood fill-based scheme while generating sparsity pattern.

In our SPION implementation, we used an embedding dimension ($D$) of size 64,
and the batch size was chosen based on the available memory size, resulting in
batch sizes of 256 for image classification, 128 for ListOps and 32 for
document retrieval. The size of convolution filter in SPION is set to
$(31\times 31)$ for all experiments. And for the threshold used in the flood
filling method, we set $\alpha$ to 96 for image classification, 98 for
ListOps, and 99 for document retrieval tasks. While performing average pooling
and upsampling, the block size for our model was determined based on the
maximum sequence length of the respective dataset. For image classification
task, we used a block size of 32 and for ListOps and document retrieval tasks,
a block size of 64 is used. Note that all the experiment results presented in
this section are averaged over 3 different executions.

###  5.1. Performance Evaluation

Table 2. Comparison of accuracy (%) on three LRA tasks

Model |  | Image
---
Classification
ListOps |  | Document
---
Retrieval
Original | 39.022 | 38.615 | 77.796
BigBird | 38.927 | 38.722 | 78.437
Reformer | 40.488 | 36.535 | 75.657
SPION-C | 39.173 | 38.889 | 76.015
SPION-F | 40.926 | 38.889 | 79.120
SPION-CF | 41.313 | 38.958 | 79.238

Convergence. Table 2 shows the accuracy of six models on three different
tasks. Our SPION-CF consistently achieved the highest accuracy in all tasks,
surpassing the highest accuracy obtained from other compared models by
+0.825%, +0.236%, and +0.801% for the three evaluation tasks. It is
interesting to see that among the SPION variants, incorporating both the
convolution filter and flood-filling scheme led to higher accuracy for all
tasks. This indicates that the convolution filter and flood-filling method
synergize with each other. More specifically, in the sparsity pattern
generation phase, the convolution filter plays a role in increasing the values
of distinct shapes appearing in $A^{s}$, thereby enhancing the ability of the
flood filling method to identify critical elements more effectively.
Additionally, we observed that the flood filling method shows more significant
effect on accuracy compared to the convolution filter. This result
demonstrates that, in sparsity pattern generation, it is more important to
consider the connectivity between elements than to simply select elements
based on their values. BigBird showed relatively lower accuracy for image
classification and relatively higher accuracy for document retrieval. This is
because the BigBird model was specifically developed to address long sequences
of text datasets.

Figure 5. Comparison of elapsed time and memory footprint per step for
training (left column), and inference time per step (right column) for three
tasks.

Speedup. Figure 5 shows the time and memory space required for training, as
well as the time required for processing inference on three tasks. Compared to
the original Transformer model, our SPION-CF achieved $1.66\times$,
$2.21\times$ and $3.08\times$ speedup per training step on each task. SPION
achieved significant speedup, particularly for tasks involving longer
sequences, such as document retrieval. As the input sequence length increases,
the number of operations in MHA increases exponentially. Consequently, in
models that handle longer sequences, the proportion of operations accounted
for MHA may be larger compared to models dealing with shorter sequences.
Therefore, in longer sequences, reducing the count of operations involved in
the MHA leads to higher performance. The figure on the right column in Figure
5 shows a breakdown of inference time of compared models. Compared to the
original Transformer, for document retrieval task, our SPION-CF achieved a
speedup of $5.54\times$ for the executing MHA operation while achieving a
speedup of $2.78\times$ in total elapsed time for inference.

Figure 6.  Breakdown of elapsed time (ms) for running MHA operations on three
tasks

To evaluate the reduction in elapsed time for each operation in the dense MHA
and sparse MHA, we compared the elapsed time of each operation in our sparse
MHA with that of the original dense MHA. Figure 6 shows a breakdown of the
elapsed time for running MHA operations in the original Transformer and SPION
on three different tasks. For the image classification task, the replacement
of the GEMM operation of $Q$ and $K$ with the SDDMM operation resulted in a
speedup of $2.55\times$. The softmax operation achieved a significant speedup
of $42.40\times$. Additionally, the SpMM operation, which replaced the GEMM
operation for $A^{s}$ and $V$, outperformed with a speedup of $2.54\times$.
This result demonstrates that the softmax function implemented in the original
Transformer model is a primary bottleneck. However, since our SPION model
performs optimized softmax by considering sparsity, the execution time for the
softmax operation is significantly reduced. Not only for the image
classification task, but also for all other tasks, our SPION achieved speedup
in every operation associated with the MHA.

Figure 7.  Comparison of training time and accuracy over different sparsity
ratio

Sparsity Ratio Comparison. Figure 7 shows the training time and classification
accuracy at each sparsity ratio on the ListOps task. Since it is not possible
to adjust the sparsity in the flood-filling scheme, we adopted the SPION-C
model to evaluate the impact of sparsity ratio. As expected, a higher sparsity
ratio leads to shorter training times, but it results in lower accuracy. Thus,
it is crucial to find the optimal sparsity ratio that balances between
achieving high accuracy and reducing training time. In the ListOps task, a
sparsity ratio of 96% provides significantly reduced training time while
maintaining a comparable accuracy. Compared to the results of a sparsity ratio
of 70%, achieving a 96% sparsity ratio resulted in a speedup of $3.26\times$,
while the accuracy only dropped by 0.903%, which is less than 1%.

Memory Footprint. As shown in Figure 5, our SPION-CF achieved a memory
footprint reduction of $4.62\times$, $7.23\times$, and $9.64\times$ compared
to the original Transformer model. And our SPION variant models resulted the
least memory footprint in all tasks. Compared to Reformer, our SPION
demonstrates comparable training time but outperforms it in terms of memory
footprint. This result shows that our SPION is able to efficiently perform
sparse MHA using sparsity patterns without requiring additional memory. In
addition, SPION-CF achieved the highest accuracy across all tasks,
demonstrating that the achieved memory savings do not come at the expense of
accuracy.

##  6\. Conclusion

Due to the compute-intensive nature of the Transformer, that must perform a
large number of MHA operations, we present a novel sparsification scheme that
leverages convolution filters and the flood-filling algorithm to dynamically
reduce the number of operations in the MHA. Our sparsification technique is
capable of identifying various aspects of the sparsity pattern in the MHA
operation and is applicable to many other Transformer models, not limited to
the encoder-only Transformer. Furthermore, we developed a parallel sparse MHA
implementation to achieve high performance. Experimental results on long
sequence datasets demonstrate that our parallel implementation on GPUs
achieves a significant performance improvement over state-of-the-art sparsity-
aware Transformer models, while maintaining comparable or better accuracy.

## References

  * (1)
  * Ainslie et al. (2020) Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. 2020.  ETC: Encoding long and structured inputs in transformers. In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_. 268–284.
  * Beltagy et al. (2020) Iz Beltagy, Matthew E Peters, and Arman Cohan. 2020.  Longformer: The long-document transformer.  _arXiv preprint arXiv:2004.05150_ (2020).
  * Child et al. (2019) Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. 2019.  Generating long sequences with sparse transformers.  _arXiv preprint arXiv:1904.10509_ (2019).
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  Bert: Pre-training of deep bidirectional transformers for language understanding. In _Proceedings of the Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_.
  * Dosovitskiy et al. (2021) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. 2021\.  An image is worth 16x16 words: Transformers for image recognition at scale. In _Proceedings of the International Conference on Learning Representations_.
  * Goldman (1990) RONALD Goldman. 1990.  Graphics gems.  _Graphics gems_ (1990), 304.
  * Kitaev et al. (2020) Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. 2020.  Reformer: The efficient transformer. In _Proceedings of the International Conference on Learning Representations_.
  * Krizhevsky et al. (2009) Alex Krizhevsky, Geoffrey Hinton, et al. 2009\.  Learning multiple layers of features from tiny images.  (2009).
  * Lagunas et al. (2021) François Lagunas, Ella Charlaix, Victor Sanh, and Alexander M Rush. 2021.  Block Pruning For Faster Transformers. In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing (EMNLP)_. 10619–10629.
  * Nangia and Bowman (2018) Nikita Nangia and Samuel R Bowman. 2018.  Listops: A diagnostic dataset for latent tree learning.  _arXiv preprint arXiv:1804.06028_ (2018).
  * Peer et al. (2022) David Peer, Sebastian Stabinger, Stefan Engl, and Antonio Rodríguez-Sánchez. 2022.  Greedy-layer pruning: Speeding up transformer models for natural language processing.  _Pattern Recognition Letters_ 157 (2022), 76–82.
  * Qiu et al. (2020) Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen-tau Yih, Sinong Wang, and Jie Tang. 2020.  Blockwise self-attention for long document understanding. In _Findings of the Association for Computational Linguistics: EMNLP 2020_. 2555–2565.
  * Radev et al. (2013) Dragomir R Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. 2013.  The ACL anthology network corpus.  _Language Resources and Evaluation_ 47 (2013), 919–944.
  * Roy et al. (2021) Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. 2021.  Efficient content-based sparse attention with routing transformers.  _Transactions of the Association for Computational Linguistics_ 9 (2021), 53–68.
  * Tay et al. (2020) Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. 2020.  Long range arena: A benchmark for efficient transformers.  _arXiv preprint arXiv:2011.04006_ (2020).
  * Tay et al. (2022) Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. 2022.  Efficient transformers: A survey.  _Comput. Surveys_ 55, 6 (2022), 1–28.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  _Advances in neural information processing systems_ 30 (2017).
  * Wang et al. (2020a) Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. 2020a.  Linformer: Self-attention with linear complexity.  _arXiv preprint arXiv:2006.04768_ (2020).
  * Wang et al. (2020b) Ziheng Wang, Jeremy Wohlwend, and Tao Lei. 2020b.  Structured Pruning of Large Language Models. In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_. 6151–6162.
  * Zaheer et al. (2020) Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. 2020\.  Big bird: Transformers for longer sequences.  _Advances in neural information processing systems_ 33 (2020), 17283–17297.
  * Zhang et al. (2021) Hang Zhang, Yeyun Gong, Yelong Shen, Weisheng Li, Jiancheng Lv, Nan Duan, and Weizhu Chen. 2021.  Poolingformer: Long document modeling with pooling attention. In _International Conference on Machine Learning_. PMLR, 12437–12446.