Transformers meet Stochastic Block Models: Attention with Data-Adaptive Sparsity and Cost
Authors: Sungjun Cho
Seonwoo Min
Jinwoo Kim
Moontae Lee
Honglak Lee
Seunghoon Hong
Venue: NeurIPS 2022 Accept (2022)
Keywords: Efficient Transformers
Sparse Attention
Stochastic Block Model

Abstract: To overcome the quadratic cost of self-attention, recent works have proposed various sparse attention modules, most of which fall under one of two groups: 1) sparse attention under a hand-crafted patterns and 2) full attention followed by a sparse variant of softmax such as $\alpha$-entmax. Unfortunately, the first group lacks adaptability to data while the second still requires quadratic cost in training. In this work, we propose SBM-Transformer, a model that resolves both problems by endowing each attention head with a mixed-membership Stochastic Block Model (SBM). Then, each attention head data-adaptively samples a bipartite graph, the adjacency of which is used as an attention mask for each input. During backpropagation, a straight-through estimator is used to flow gradients beyond the discrete sampling step and adjust the probabilities of sampled edges based on the predictive loss. The forward and backward cost are thus linear to the number of edges, which each attention head can also choose flexibly based on the input. By assessing the distribution of graphs, we theoretically show that SBM-Transformer is a universal approximator for arbitrary sequence-to-sequence functions in expectation. Empirical evaluations under the LRA and GLUE benchmarks demonstrate that our model outperforms previous efficient variants as well as the original Transformer with full attention. Our implementation can be found in https://github.com/sc782/SBM-Transformer.

Transformers meet Stochastic Block Models: Attention with Data-Adaptive Sparsity and Cost
Sungjun Cho
Seonwoo Min'
Jinwoo Kim
Moontac Lee
Honglak Lee"
Scunghoon Hong
LG Al Research
2 KAIST
University of Illinois Chicago
Abstract
To overcome the quadratic cost of self-attention, recent works have proposed various sparse attention modules, most of which fall under one of two groups: 1) sparse attention under a hand-crafted patterns and 2) full attention followed by a sparse variant of softmax such as a-entmax. Unfortunately, the first group lacks adaptability to data while the second still requires quadratic cost in train- ing. In this work, we propose problems by endowing each attention head with a mixed- membership Stochastic Block Model (SBM). Then, each attention head data-adaptively samples a bipar- tite graph, the adjacency of which is used as an attention mask for each input. During backpropagation, straight-through estimator is used to flow gradients beyond the discrete sampling step and adjust the probabilities of sampled edges based on the predictive loss. The forward and backward cost are thus linear the number of edges, which each attention head can also choose flexibly based the input. By assessing the distribution of graphs, we theoretically show that SBM-Transfornee is a universal approximator for arbitrary sequence-to-sequee functions in expectation. Empirical evaluations under the LRA and GLUE bench- marks demonstrate that our model outperforms previous efficient variants as well the original Transformer with full attention. Our implementation can be found SBM- in bnttt:://// //i thub. com/sc http:///////....aanntootoooe Transformer
model that resolves both
Introduction
The Transformer
S architecture has been the go-to method for encoding sequential data, due to its superior performance various tasks such machine translation (28), image classification [14]. and protein language modeling 32]. Its key strength stems from the multi-head attention module, where so-called attention score matrix computes how contextually important one token another for all possible token pairs. Each Transformer layer simultancously pools the token representations based on the attention scores, eventually returning contextualized features without sequentially traversing through the input sequence as its recurrent neural network-basedr predecessors [16].
well-known drawback of the original Transformer is its high computational cost in time and memory that increases quadratically with sequence length. This is due to the full pairwise computation of attention scores, which prohibits applying it in tasks involving long-range dependencies such as document summarization [17] or high-resolution image processing [48] Many works have thus focused on developing more efficient alternatives by exploiting fixed learnable attention sparsity
patterns g 46 20 [2], low-rank approximations 44 433, or kernelized attention modules [19 9]
Even though the efficient alternatives hold theoretical expressibility guarantees 45 they are far from sufficient, still failing to convince practitioners to replace the original Transformer. We believe this is mostly due their lack of adaptability. They apply the same modifications unanimously sparsify all the attention modules across layers, without considering the tasks at hand. Such strategy
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Linear
M
Concatenate
Implicitly compute
SBM
edgelist
QKT
Query Key K
OM
O
Scaled Dot-Product Attention
VD
et
Linear
Linear
Linear
The attention module SBM-Transformer. In multi-head attention, each attention head samples a bipartite graph connecting queries to keys from an underlying SBM. The adjacency of the sampled graph is used as an attention mask to compute the dot products only for the sampled edges.
Figure 1:
imposes inductive bias too strongly and often leads sub-optimal cost vs. performance trade-offs in downstream tasks [27]. In this work, we argue that to retain the utmost potential of Transformers each attention module should have the ability to flexibly choose between sparse and full attention. This is especially evident when considering many state-of-the-art systems suggest the need for a mixture of dense and sparse attention layers. For example, qualitative analysis on pretrained BERT showed that lower layers exhibit broad dense attention while upper layers perform focused sparse attention [10]. In the case of GPT-3 [6. the Transformer blocks are manually arranged to alternate between dense and sparse attention.
To contribute to the efficient Transformers lineage, we propose 3BM-Transfomen capable of adjusting its attention sparsity data-adaptively based without fully computing the attention score matrix (Figure [T. Leveraging mixed- mexxeermmmbrrship Stochastic Block Model (SBM) (2, each attention head samples bipartite graph connecting queries to keys. Then, the adjacency of the sampled graph used an attention mask so that only attention scores corresponding sampled edges are computed. The overall computational cost is linear in the number of edges, which can range from linear to quadratic in sequence length depending on the data and task under concern. Each attention head is equipped with its own underlying SBM, enabling the model to diversify the attention sparsity across heads and layers. By incorporating a straight- through estimator 4 in the discrete graph-sampling step, SBM- Tra--rrasformrr enjoys end-to-end differentiability and can find the proper attention sparsity based solely upon minimizing the predictive loss. The model can also easily be further regularized by penalizing the number of sampled edges, which results in lighter model using less computational resources during inference. To the best of our knowledge, our method is the first Transformer architecture that can data-adaptively choose between linear to full attention with respective computational costs. To summarize, our main contributions are as follows:
We present TrM--rrasformer, novel Transformer of which each attention head can adaptively adjust its attention sparsity as well as computational cost based on the input data.
To demonstrate the benefit of this flexibility, we theoretically prove that Tranrornsform retains universal approximability, and also stress-test the model under synthetic task where full attention required to achieve 100% accuracy.
Evaluations LRA and GLUE benchmarks show that SBM-Transformer outperforms pre- vious efficient Transformer models as well as the vanilla Transformer with dense attention.
2 Related Work
In this section we discuss previous efficient Transformer variants and several works similar to ours with respect to adaptively learning sparse attention patterns. We also review several works on SBMs.
Efficient Transformers. Many efficient Transformers tackle to reduce the quadratic cost of multi- head attention with different approaches. While we discuss only a handful of representative ap- proaches, much more comprehensive survey can be found in [37]. The Linear Transformer [19] achieves linear complexity by replacing the softmax with a low-rank kernelized function. Lin- former [40] and Nyströmformer 43 use a similar approach by low-rank approximating the attention score matrix. Performer [9] uses positive orthogonal random features to approximate the softmax kernel. Reformer 20 gathers similar tokens together through Docality-sensitive hashing (L.SH) and performs attention amongst tokens within the same bucket. Of all methods above, our method is

---

most similar to Reformer, in the sense that we adaptively assign queries and keys into clusters and form low-rank sparse attention pattern. However, our method performs soft-clustering with much less structural constraints, allowing each attention head to represent a wider variety of dependency structure and to adjust its sparsity towards full attention if needed.
Adaptive Sparsity. With respect to flexible training between sparse and dense attention, there exist some works that parameterize how sparse the attention pattern should be based on the input. The Adaptive Sparse Transformer [T] proposed replacing the usual softmax activation with a-entmax, in which the parameter can be differentiably trained to adjust the activation between softmax and sparsemax activation (25). SparseBERT (4] uses a differentiable masking technique where each attention mask is sampled from a sigmodl distribution using data- catt-inndeeedeen mask probability parameters. While these methods possess the flexibility to adjust between sparse and full attention based on data, they still require full computation of the attention score matrix before sparsification, and hence are unable to leverage the learned sparsity towards better model efficiency. To the best of our knowledge, ours is the first work to he able adaptively tune its attention sparsity between sparse to full attention without requiring the explicit computation of the attention score matrix, thereby avoiding quadratic cost when possible.
Stochastic Block Models. The Stochastic Block Model (SBM) a generative model that encodes the latent structure of graphs by grouping nodes into clusters. By modeling the cluster-membership of each node as well as inter-cluster relationships, SBMs can represent a wide variety of graph structures, which is a feature especially useful for generating new graphs or predicting missing edges in noisy data [.] The standard SBM assigns each node a single cluster, and the probability of an edge between two nodes strictly depends on the corresponding clusters. Several structural extensions include overlapping SBM [2 and menhershehhr SBM (2). which allow each node to be assigned to multiple clusters. The underlying SBM used by our framework mostly resembles these two variants, while the edge probability is modeled by a nonlinear function of two node embeddings rather than a bilinear one. There exist many other extensions including degree-correcter SBM [18] for guati- and hicrarchical SBM [29] for multiplex- multiplex-graphs. Further details can be found in a recent survey [55.
3 Preliminaries: Sparse Transformers
We first introduce the full attention mechanism used in the original Transformer as well as masked attention which will serve backbone of our approach.
3.1 Full Attention
In vanilla Transformer 38 each attention head takes sequence of token features as input X E gaxd where n is the sequence length and d the embedding dimension. Weight parameters Wq, W E Rd and W Rddxd with head-dimension d, first maps the input features X into query Q. key K, and value V. respectively. Then, the attention score matrix is computed with scaled dot-product of queries and keys followed by row-wise softmax activation o a-)). Note that explicit computation of this matrix is the main bottleneck of full attention, incurring O(n²) asymptotic cost in both time and memory. The value features V are then pooled based on the attention scores, returning the output token representations. Altogether, the operation performed by each attention head can be written as
Q XWO K = =XW V VV
Attn (x) (QK") V.
=
(2)
3.2 Masked Attention
One way to remove the quadratic bottleneck from the attention score matrix is to apply a binary mask M E {0, 1 and compute the scaled dot-products K,K5/ / /dnh only if M.;- = 1. In presence of an attention mask, the operation is modified to
QKT
AttDmask XX M) OM
M A
V
(3)
Vdn
exp(A.g)
if Mij =1
TM (A)i =
ERE(K(M exp( Aax)
(4)
0
otherwise
3

---

Q.K
Forward
+---.
S
Backward
Clusters
Attention Mask
SBM + STE
Output
Attn mask
X
Imputs
K
0L/0X
Query/Key
Nodes
V
Figure 2: An illustration of the attention mechanism in queries and keys to the node representation space through a shared MLP. The graph sampling module samples attention mask from a Stochastic Block Model (SBM) parameterized by the node and cluster embeddings. The discrete sampling step is differentiable via Etraigh-Thhoouh Estimator (STE). Given the mask, the output is computed via masked attention.
Each head first maps
where indicates entry-wise multiplication. Note that the masked-soffmar dM operator only computes unmasked terms, ensuring that each (i, j-th attention score survives as nonzero if and only if Mij 1. This is thus equivalent filling in the (i. (-th attention score with -00 if Mij then applying the standard softmax operator. Most sparsity-based efficient Transformers fall under this formulation, while using different methods to either manually fix or learn the mask M. For instance, local attention
3 46 with sliding window sets Mij = 1 if J - -j < c for some context window size e while Reformer (20 sets Mij- if Q; and K2 are hashed into the same bucket.
8
4
Our Method: SBM-
Here discuss the details of SBM-Transformer (Figure 2 We first illustrate the forward step of our attention module and how the underlying SBM 2 of each head, from which we sample our attention masks, parameterized by the input tensors. We then discuss how the model enables end-to-end differentiability despite the discrete graph sampling step.
4.1 Forward step with the Stochastic Block Model
In our framework, we view the attention mask M as an adjacency matrix of a bipartite graph that connects queries keys, and let each attention head sample an adjacency matrix that best represents the contextual dependencies amongst input tokens. order to efficiently sample adjacency matrices while avoiding the quadratic cost, the distribution of graphs must first be parameterized with sub-quadratic number of latent variables. Stochastic Block Models fit perfectly for our purpose as it models graphs that low-rank structured with k latent clusters, allowing full parameterization using O (nk) memory. More concretely, the SBM distribution is defined by two nonnegative node- to-cluster memberships Y, Z E R and a so-called block matrix B E R that stores the inter-cluster connection probabilities. The probability of node being connected node j is computed as p(i.) j) Y,BZ,. Equivalently, the expectation of the adjacency matrix sampled from A ~ SBM(Y, B, Z) can be written as E[A] YBZ..
For proper parameterization of the SBM, we must infer the nonnegative node- me--ermmmbbessssse and block matrix from the queries and keys. To do so, equip each attention head 2-layer MLPd,-vd, with ReLU activation, and a set of k trainable cluster-ens enaeterrembediings C E & First, our model computes the block matrix S E x by taking dot products amongst duuster-ebeeddings C' followed by 2-dimensional softmax activation. The node embeddings are obtained by processing each query and key through the MLP -dd, mapping token representations into the node representation space. The memberships of query and key nodes, which we denote by Q and K, are then inferred by taking dot products of node and cluster embeddings, followed by sigmoid function. The block matrix S. query node-memberships Q, and key node-memberships K altogether provide well-defined parameterization for the SBM. Thus, bipartite graph adjacency M E {0, 1)axm can be sampled from M SBM(Q, S, K) with expectation E. E QŠKT: the probability of connecting query Q, key K, equals p(i.j) Q.KK? Formally, the sampling procedure be written

---

Algorithm 1: fastRG((, B, 3)
Input :Y E RTT* B Ri** Z IR
O {0, 1)mmm with E[M] YBZ
1 Compute diagonal matrices Dy = (diag(1Y)) and Dz = (diag( ( 12)
Column-normaile Y =YDD' and Z ZDz
3 Compute = BDZ
4 Sample number of edges m ~ Poisson (1Bson Initialize M
6 for i = 1: m do
Sample (U, V) from [l....,k k3 x {I,..., k} with Pr(U u, =t) Ban
8
Sample source I from (I,...,n) with Pr(1- = Yiu.
Sample destination J from [l,...,n] with Pr(J -V
10
Set MI = 1.
11 end
Š softmax(CC*)
(5)
Q sigmoid(MLPP) -dd, (Q)CT)
(6)
K sigmoid(MLP... da -d, (K)C²)
(7)
M SBM
(8)
For the last sampling step, we incorporate a fast random graph sampling algorithm fastRG (Alg. I 33) that can sample graphs from a SBM time and memory asymptotically linear the number of edges. One advantage of fastRG is that each edge can be sampled in parallel, allowing high efficiency with the help of multiprocessing. A more significant feature of the method that the number of edges, which determines the overall cost, is sampled from Poisson distribution with input-dependent mean (Line 4). Thus, the model can dynamically adjust computational cost between linear and quadratic in sequence length based on the data.
Figure shows example placements of nodes and clusters on the d, -dimensional space show how the sparse structure is determined. If all nodes and clusters are gathered closely, then all entries in Q and K become close to resulting in p(i.j) well-separated but each surrounded by some set of nodes, S becomes close to diagonal while each row in Q and K is close to a one-hot hot vector indicating the cluster nearby. Such setting leads to block diagonal mask similar to LSH bucketing of Reformer [20]. Lastly, if all clusters are far apart from the nodes, both Q and K approximately equal zero, zeroing out all the edge probabilities.
for all i.j and hence dense M. If clusters are
4.2 Backward Step with Straight-Throug Estimator
The graph sampling procedure is naturally discrete operation. Thus, naive backpropagation cannot learn the proper parameterization for the SBM that minimizes the predictive loss. To cope with this non-differentiabiliyy.
incorporate Straight- Through Estimator (STE) [4] to pass the gradient beyond the discrete sampling step. The STE enables providing the gradient 0C/0M; to the probability for each sampled edge (i.j) (Eqn. 99. It works as if we had used continuous mask each sampled edge instead of the binary mask M during
M @ M that stores the probability forward propagation. This way, the probabilities of sampled edges be learned end-to-end: the gradients provide information on whether each sampled edge was useful or not for prediction.
JC
Q.K Vdh
OC
if M =1 where
M O QKT
DAy
Op:
aM
Vdh
otherwise
5

---

(a) Dense
(b) Block-diagonal
(c) Sparse
Query-nodes
Key-nodes
Clusters
E M
Representative examples from the SBM and resulting mask expectations (darker grid indicates edge probability closer 1). (a) The expected mask is dense if all nodes and clusters are collapsed within small region. (b) Clear-cut groups in the embedding space induce a block-diagonal mask. (c) Clusters located far apart from nodes lead to sparse masks.
Figure 3:
Random Edge Exploration. While this approach enables backpropagation in the same (m)) cost as in the forward step, this comes the expense of not being able to propagate information through edges that were not sampled. This can be problematic when edge probability accidentally collapses to zero, after which the edge becomes unlikely to ever be sampled even when it may be useful for prediction task at hand. Therefore, we add small perturbation å >> 0 to each edge probability Pij. allowing the model to explore new edges and resuscitate their sampling probabilities if necessary. We find that 5 as small as 0.01 significantly helps in practice, and thus use this edge exploration scheme during training for our experiments.
Wouldn't 't the model always prefer full attention? Note that the gradient @C/0p.j can be positive, which suppresses the probability of edge (i. j). At first, may seem counter-intuitive why the model would ever limit itself to using fewer edges during training without any sparsity- -bsed regularizations. One explanation is that masked attention provides head dimensions. Under full attention, it is known that the representational space of attention score matrices is limited by the head dimension and softmax activation (5). This limitation inevitably introduces unwanted noise in the attention scores especially when working with long sequences. In however, the structural sparsity in masked attention introduces another dimension that induces a larger space of row-stochastic matrices (full attention is a special case of masked attention where Mi; for all i, Therefore, it is reasonable that the model may encourage sparsity to leverage the additional expressiveness assuming the loss landscape has local optima within the sparse attention regime. Our experiments on the LRA benchmark show that this is indeed the case, as our Tra-Transfrrmer converges to an average attention sparsity of 20% to 30% while outperforming Transformer with full attention. We also show in the experiment that we can easily incorporate additional regularization that further encourages sparse attention masks.
easy way to reduce attention scores under finite
4.3 SBM-Transformer is Universal Approximator
Leveraging previous work the theoretical expressiveness of sparse attention 45
46
show that
SBM-Transformer with a small modification Dodifiicatiorrretinns retains the same level of expressibility as full attention. Specifically, we show that the low-rank structure of the underlying SBMs does not degrade the expressive power of Transformer, and that SBM-Transformer can universally approximate arbitrary functions with (nn) connections. For brevity, we provide rough overview the proof and defer further details to Appendix A.
Theorem 1. Let f E F be class of continuous eequence-to-sequeeee functions. TSBM denote the class of SBM-Transformers with h attention heads, head dimension, and dimensions in hidden layers. Then for any 0 and 1 S p oc, there exists function g TSBM such that
J [1((X) [ E[g Se
(10)
Herr we consider a variant of SBM- Transformer where self- self-loops are added manually (i.e. M,, for all i). While this is useful in theoretical analysis, we find that not having self self-loops slightly helps in empirical performance and hence omit self-loop for the main experiments.

---

1.0
0.6
Full-atth
Reformer
Training Loss
Mack Denetty
Performer
0.7
Nystrom
0.6
Lincarized
0.3
Linorrer
SBM(urs)
0.3
0.0
250
1500
0.2
1000 1500 2000 Epoch
Epoch
Figure 4: Loss (left) and mask density (right) of SBM-Transffrmee during training the synthetic task. 3BM-Transformer successfully converges to zero loss by tuning itself towards full attention.
According the main theorem of Yun et al. (2020) 44], SBM-Transformer achieves universal approximability if I) each node attends to itself, 2) the aggregation all attention patterns contains Hamiltonian path, and there exists path between all node pairs. While the first condition is trivially true due to our modification, the other two conditions require careful choice of three SBMs. Here first parameterize one SBM hard-assign tokens into equally-sized clusters, inducing block-diagonal attention pattern. The other two SBMs are parameterized such that the two graphs together form a star graph with i global relay tokens. Combining the three attention patterns lead to that satisfies all three conditions, hence proving the theorem.
parameterization of
Experiments
For empirical evaluations, we first use a synthetic task to show that our model is flexible enough to learn towards full attention when needed in contrast previous works. We then experiment Long Range Arena (LRA) [36]. benchmark widely used to assess the capacity of efficient Transformers learning long-range contexts across different modalities. Lastly, we show results on the GLUE benchmark [39 to assess the performance of SBM- Tranorrnsfrre in downstream NLP setting. All experiments were run on a remote GCP server equipped with 16 NVIDIA A100 Tensor Core GPUs.
5.1 Synthetic Task: Finding Repeated Tokens
Dataset. We formulate a token-level binary classification task as follows: each input sequence consists of N integers, each of which is uniformly sampled from 2,, N.). We use N 256 in our setup. The prediction target is sequence of equal length, where each token is labeled if there exists duplicate somewhere within the sequence, and 0 otherwise. Below is simple example with N that illustrates the task. We measure the performance models via binary cross-entropy loss.
N).
Input: 1 4 3
Target:
Methods. For this task, we compare SRMM-Trasformer with = 128 clusters against various efficient Transformers: Linear Transformer [9 Linformer [4). Reformer 220, Performer (9. and Nyströmformer [43]. Across all methods, we use single-layer and single-head architecture with 32 hidden dimensions. Note that due to this constrained setting, the sole head must perform full attention to compare each token to all the others order attain 100% accuracy. All models are trained for 2000 epochs where new batch of sequences is sampled on-the-fly each epoch. We use batch size of 256 and learning rate of 1e-3.
Results. Figure shows the training loss curves of each baseline method as well as SBM- Transformer. Full attention quickly converges to 100% accuracy, which expected as computes all possible pairwise interactions by default. Other models that apply low-rank kernelized attention fail to achieve the same level of accuracy, due to limited expressibility under the constrained setting. Though
converges more slowly compared to full-attention, it demonstrates the ability to drive itself towards full-attention, eventually attaining zero loss.
7

---

Model
LiSTÖps(2K)
TEXT(3K)
RETRIEVAL(4K)
IMAGE(IK)
PATHUFINDER(K)
Avg.
Full-attention
37.22
64.93
79.55
40.38
74.26
59.27
Linearized 19
37.46
64.90
81.10
38.48
74.61
59.31
Reformer 20
22.92
64,70
77.25
43.63
70.28
$5.76
Performer
9
18.25
65.00
79.01
70.79
54.57
Linformer
40
38.44
56.28
78.09
39.53
67.62
55.99
Nyströmformer 43
37.22
65.46
79.35
43.02
71.97
59.41
SBM Transformer (ours)
37.445 (20.09%%
65.79 (26.10%)
(29.46%)
41.31 (20.49%)
75.12 (18.56%))
59.93
Table 1: LRA benchmark results. The sequence lengths are shown
each task. For SBM- Transformer, we report the average attention sparsity across all layers and heads during test time in parentheses. Bold and underlined results indicate best and 2nd best test accuracy for each task.
A
(22T
TEXT(3K)
RETRIEVAL(4K)
IMAGE(IK)
PATHFINDER(IK)
Avg.
0
37.45 (20.09%)
65.79 (26.10%)
80.00 (29.46%)
41.31 (20.49%)
75.12 (18.56))
59.93
10-4
37.76 (10.48%)
65.48 (26.26%)
79.93 (24.62%)
41.35 (10.70%)
75.46 (5.16%)
60.00
10-3
38.23 (10.46%)
65.18 (26.03%)
80.00 (21.70%)
41.17 (24.60%)
74.49 (3.82%)
59.81
10-2
38.2 (2.95%)
65.59 (22. (22.43%)
80.44 (6.99%)
42.20 (3.95%)
72.79 (3.76%)
59.84
10
37.76 (1.15%)
64.48 (10.62%)
79.46 (2.49%)
41.35 (1.33%)
73.79 (2.61%)
59.37
Table 2: LRA results of Tran-fooorormer with increasing sparsity regularization weight A. Bold results indicate best accuracy for each task and percentage parentheses indicate average attention density. Sparsity regularization helps in reducing computational cost with small drop in performance.
5.2 Long Range Arena (LRA)
To demonstrate that the flexible inductive bias of SBM-Transformer effective for modeling long- range dependencies, test SBM-Transfomme Transformer against previous work the LRA benchmark. We also test how the performance is affected with respect to applying a sparsity- based regularizer.
Dataset. LRA (36] consists of five different testbeds with varying modalities: LISTOPS [26 is 10-way classification task to map a sequence of single-digit numbers and different set operations, to its corresponding solution. TEXT 24 is a binary classification task where byte-level IMDB movie reviews must be classified into one of positive
negative sentiments. RETRIEVAL 30 also a char-level binary classification task, where two sequences from ACL Anthology papers are given as input, and the model must predict whether there exists a citation link between them. IMAGE [21] is a 10-way classification task mapping flattened pixel-sequences from CIFAR-IO to its class. PATIIFNDDR 23 provides flattened pixel-sequences from image and the model must decide whether two circles the image are connected by dashed line. For this benchmark, we use the Py PoTcch implementation of LRA provided by the authors of Nyströmformer [4] and adhere to the same train-test splits. Performance in all five tasks is measured using classification accuracy.
Methods. We compare SBM-Transforee against the same baselines with the synthetic task above. For fair comparison, we set all Transformer models to use the default setting used in 43 which fixes 2 layers, 2 attention heads, and 64 embedding dimensions. For SBM Transformer, we use &o = 128 clusters. The output token representations are mean-pooke ppoled to obtain the sequence representation for all tasks. More details on the architecture setups can be found Appendix C.
achieves the best overall performance, ranking first in two tasks, and second in one other. SBM- Transformer also outperforms full attention in all five tasks while computing 30% less attention scores average, which supports our claim that masked attention with partial attention score computations can be preferred over full attention depending find that flexibility of SBM is indeed beneficial, as Reformer struggles in LISTOPS, most likely due the inability of block-diagonal masks model hierarchical contexts.
Results. Table[] the test accuracies of each method. Our
the task. With respect the attention mask structure,
Mask Density Regularization. To test if the model can effectively learn under constraint on the computational cost, we also test the model under a sparsity- regularizer that discourages excessive use of query-key edges. We penalize each sampled edge by adding to the predictive

---

Relative FLOP Count
Relative Peak Memory Usage
Model
L(2K)
T(3K)
R(4K)
I(K))
P(IK)
L(2K)
T(3K)
R(4K)
I(IK)
P(IK)
Full-attention
38
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
Lincarized 19
0.02
0.01
0.02
0.04
0.04
0.18
0.16
0.12
0.42
0.42
Reformer 20
0.05
0.03
0.05
0.10
0.10
0.39
0.31
0.18
0.72
0.72
Performer 9
0.18
0.12
0.18
0.36
0.36
0.76
0.70
0.60
0.96
0.96
Linformer 40
0.33
0.22
0.33
0.66
0.66
0.26
0.22
0.14
0.34
0.34
Nyströmformer [3
1.09
0.70
1.09
2.37
2.37
0.34
0.27
0.16
0,70
0,70
SBM-Transfone Transformer (ours)
0.07
0.23
0.08
0.27
0.29
0.19
1.01
0.19
0.39
0.48
Table 3: Per-example relative FLOP count and peak memory usage during LRA inference.
Ops
Image
Retrieval
Pathfinder32
Text
40
Mask deusity (%)
30
20
(no reg.)
(reg.)
L1 12
L1
12
L1 t2
11
L1
i 12 11
L1 11
Figure 5: Average and standard deviation of density of masks sampled across the test set for each LRA task. The æ-axis indicates the lower (LI) and upper (L2) layers and each bar represents the density averaged between the attention heads in each layer.
loss weighted regularization term ACs, where C. denotes the average mask density across all attention heads. weights. Under strong regularization, the model surprisingly retains competitive performance while significantly reducing the average mask density. This indicates that similar local optima are shared across regimes with varying attention density in the loss landscape, and the regularization term is able to drive the model towards finding optimal attention scores with smaller density.
the performance of SBM- Transformer across varying regularization
Efficiency. Furthermore, we compare computational costs during inference by measuring FLOP count and peak memory usage. For Due to lack of support for sparse tensor operations in existing FLOP-counters. we measure FLOP counts by manually enumerating through each tensor operation. Table shows that TrM--rannformer comparably efficient across all tasks except for TEXT, where 1BM--ransformer showed the largest average mask density. Note that while the cost of other baselines are fixed after initialization, the cost of -ran-forrformer is data-adaptive and can vary upput-by-input Further analysis and qualitative examples demonstrating the input-dependent attention mask densities can be found in Appendix C.
we test the model trained under À = 10-1
Layerwise Diversity in Sparsity. We also compare the densities of masks sampled at each layer of SBM-Transformer during test time to examine whether our model is capable of diversifying sparsity across layers for better performance. Recall that this allows models to gather information different levels, seen in pretrained BERT where lower layers focus on the overall content via dense attention while upper layers gather syntactic information with tree-like patterns [T For each of the five tasks, we pick two nighest-performing models (one for unregularized and another for regularized) for measurement. Figure shows the average layer-wise mask densities of unregularized and regularized SSM-Transformers across different tasks. We find that under regularization, the two layers can differ by more than 10% in tasks such as LISTOPS and IMAGE. This may be due to hierarchical and compositional structure of the two tasks. We also find that the variation is relatively low in Text with densities around 25%, indicating that the task requires broad attention overall. Lastly, the standard deviation is extremely large in upper layers for PATHFINDER, showing that it samples wide variety masks depending on the input.

---

5.3 General Language Understanding Evaluation (GLUE)
To check whether its strong performance demonstrated in LRA extends to the downstream NLP setting as well, evaluate SBM-Transformer against baselines on the GLUE benchmark [39]
Dataset. We consider four NLP tasks in GLUE 39 SST-2 35 consists of movie reviews the model must predict their positive or negative sentiments. For QQP 7. the task is to determine whether one question is paraphrase of the other given pair of questions. MNLI 22 sentence pairs, each with a target label indicating whether the two sentences are connected through entailment, contradiction, or neither. QNLI 31 consists of sentence-question pairs and the task is to determine whether the sentence contains an answer to the question. Each task is formulated as sequence classification, and we measure performance by Fl score the respective validation sets.
consists of
Methods. Following previous work 443 we arrange small variant of BERT [3 with 4 layers, attention heads, and 512 embedding dimensions. We replace full attention with each attention module used in previous experiments. For TSM--ransformer, we use 128 clusters without sparsity regularization (i.e. À = 0). Here, we find that adding local attention significantly boosts performance, and thus fix a sliding window of size 64 to SBM Transformer. We first pretrain each model under the masked language modeling objective for 50 epochs on corpus with text from English Wikipedia, BookCorpus [50]. and RealNews 47). We then finetune each pretrained model for 5 epochs on the GLUE training sets. More details on the architecture and training setup can be found in Appendix C.
Results. Table reports the Fl scores of each method on different NLP tasks. SBM- Transformer performs competitively against full attention overall, and outperforms all baselines in SST-2 and QQP. We also find that the fine-tuned SBM-Transien Transforfrre models use 13.5% dense attention masks on average across tasks, showing that the model can encode use- ful information from input sentences effectively under highly sparse attention.
Model
SST-2
QQP
MNLI
QNLI
Full-attention
38
89.8
84.7
84.0
85.0
Reformer
20
89.3
84.4
83,9
84.0
Performer
9
82.0
65.6
71.4
59.3
Linformer 40
82.0
83.2
79.3
82.5
Nyströmforner +
89.7
83,2
84.1
84.9
SBM- Transformer (ours)
89.8
85.2
83.5
83.6
Table 4: GLUE benchmark results. Bold results indicate best accuracy for each task.
6 Conclusion
We propose SBM-Transformer sparsity between sparse and full attention without the need to explicitly compute the full attention score matrix. Theoretically, we show that our model enjoys the same expresiibility as the original Transformer due to the flexibility of the latent SBM. Empirical experiments on LRA and GLUE show that our model performs competitively against previous of efficient Transformers.
efficient Transformer that data-adaptively choose its attention
Nonetheless, there are limitations due to sparse tensor operations being less optimized on GPU kernels. In the LRA experiments, found that dense counterparts while its memory usage is much lower. While previous sparsity- attention mechanisms with block-sparse attention much more amenable for GPU computation 446 g 3
result in longer runtimes compared to
our work requires an architecture with better workload balancing and acceleration under unstructured sparsity, for which there ongoing work [41] 49
We still believe this work is valuable as it is the first approach to induce per-example attention sparsity, allowing the model to adjust its computational cost based the input. The cost being dependent on the number of edges also allows practitioners to easily impose constraints based on the available computational resources. We hope to see more GPU-friendly tensor operations optimized for fine- grained sparsity in the future, at which point the value of this work will increase even further. As we propose foundational replacement for the scaled dot-product attention module in the Transformer architecture, we do not expect any immediate negative societal impact due to this work.
Acknowledgments and Disclosure of Funding
We would like to thank Kun Dong for the insightful comments. This work was supported by Institute of Information & communications Technology Planning Evaluation (ITPP) grant funded by the Korea govermment (MSIT) (No. 2022-0-00926, 2022-0-00959, 2021-0-02068, and 2019-0-00075.
10

---

References
[1] E. Abbe. Community detection and stochastic block models: recent developments. The Journal of Machine Learning Research, 18(1):6446-6511 2017.
[2] E. M. Airoldi, D. Blei, S. Fienberg, and E. Xing. Mixed membership stochastic blockmodels. Advances in neural information processing systems, 21, 2008.
[3] I. Beltagy, M. E. Peters, and A. Cohan. Longformer: The long-document transformer. CoRR, abs/2004.051525 2020.
[4] Y. Bengio, N. Léonard, and A. C. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. CoRR, 100./330,.432, 2013.
[5] S. Bhojanapalli, C. Yun, A. S. Rawat, S. J. Reddi, and S. Kumar. Low-rank bottleneck in multi-head attention models. CoRR, abs/2002.07028, 2020.
[6] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, Shyam, G. Sastry, A. Askell, Agarwal, A. Herbert- verbert-Voss,, G. Krueger. T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners. CoRR, abs/2005. 14165,,,,,1,,,, 2020.
[7] Z. Chen, H. Zhang, X. Zhang, and L. Zhao. Quora question pairs. 2017.
[8] R. Child, S. Gray, A. Radford, and I. Sutskever. Generating long sequences with sparse transformers. CoRR, abs/1904. 10509, 2019.
[9] K. Choromanski, V. Lihhosshertoo, D. Dohan, X. Song, A. Gane, T. Sarlós, P. Hawkins, J. Davis,
A. Mohiuddin, I.. Kaiser, D. Belanger, J. Colwell, and A. Weller. Rethinking attention with performers. CORR, abbs2200..44799 2020.
[10] K. Clark, U. Khandelwal, O. Levy, and C. D. Manning. What does BERT look at? an analysis of bert's attention. CoRR, abs/ 1964.0431, 2019.
[11] G. M. Correia, V. Niculae, and A. F. T. Martins. Adaptively sparse transformers. CoRR, abs/1909.00015. 2019.
[12] G. Daras, N. Kitaev, A. Odena, and A. G. Dimakis. SMYRF: efficient attention using asymmetric clustering. CoRR, abs/2010.05315, 2020.
[13] Devlin, M.-W. Chang. K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[14] A. Dosovitskiy, L. Beyer, Kolesnikov, D. Weissenborn, Zhai, T. Unterthiner, M. Dehghani. M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An image is worth 16x words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[15] T. Funke and T. Becker. Stochastic block models: A comparison of variants and inference methods. PLOS ONE, 14(4):1-40, 04 2019.
[16] S. Hochreiter and J. Schmidhuber. Long short-term memory. Neural computation, 987:1733-- 1780, 1997.
[17] I.. Huang, S. Cao, N. Parulian, H. Ji, and L. Wang. Efficient attentions for long document summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1419-1436. 2021.
[18] B. Karrer and M. E. Newman. Stochastic blockmodels and community structure in networks. Physical review E, 83(1):016107, 2011.

---

[19] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. CoRR, abs/2006.16236 2020.
[20] N. Kitaev, L. Kaiser, and A. Levskaya. abs/2001.04451, 2020.
Reformer: The efficient transformer.
CoRR,
[21] A. Krizhevsky. Learning multiple layers of features from tiny images. 2009.
[22] P. Latouche, E. Birmelé, and C. Ambroise. Overlapping stochastic block models with application to the french political blogosphere. The Annals of Applied Statistics, 5(1):309-336, 2011.
[23] D. Linsley, J. Kim, V. Veerabadran, and T. Serre. Learning long-range spatial dependencies with horizontal gated-recurrent units. CoRR, 185.0883 2018.
[24] A. L. Maas, R. E. Daly, P. T. Pham, D. Huang, A. Y. Ng, and C. Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pages 142-150, Portland, Oregon, USA, June 201 Association for Computational Linguistics.
[25] A. F. T. Martins, M. V. Treviso, A. Farinhas, V. Niculae, M. A. T. Figueiredo, and P. M. Q. Aguiar. Sparse and continuous attention mechanisms. CoRR, abs/2006.07214, 2020.
[26] N. Nangia and S. R. Bowman. Listops: A diagnostic dataset for latent tree learning. CoRR, abs/1 0bs//884.000/ 06020, 2018.
[27] S. Narang, W. Chung, Tay, W. Fedus, T. Févry, M. Matena, Malkan, N. Fiedel, N. Shazeer, Z. Lan, Y. Zhou, W. Li, N. Ding, J. Marcus, A. Roberts, and C. Raffel. Do trans- former modifications transfer across implementations and applications? CoRR, abs/2102.11972. 2021.
(28) M. Ott, S. Edunov, D. Grangier, and M. Auli. Scaling neural machine translation. In Proceedings of the Third Conference on Machine Translation: Research Papers, pages 1-9, Brussels, Belgium, Oct. 2018. Association for Computational Linguistics.
[29] T. Peixoto. Hierarchical block structures and high-resolution model selection in large networks. Physical Review X, 4(1):01 2014.
[30] D. R. Radev, P. Muthukrishnan, and V. Qazvinian. The ACL Anthology network. In Proceedings of the 2009 Workshop on Text and Citation Analysis for Scholarly Digital Libraries (NLPIR4DL), pages 54-61, Suntec City, Singapore, Aug. 2009. Association for Computational Linguistics.
[31] P. Rajpurkar, J. Zhang, K. Lopyrev, and P. Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 2383-2392, Austin, Texas, Nov. 2016. Association for Computational Linguistics.
[32] R. M. Rao, J. Liu, R. Verkuil, J. Meier, J. Canny, P. Abbeel, T. Sercu, and A. Rives. Msa transformer. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8844- 8856. PMLR, 18 24 Jul 2021.
[33] K. Rohe, J. Tao, X. Han, and N. Binkiewicz. A note on quickly sampling a sparse matrix with low rank expectation. The Journal of Machine Learning Research, 1911:3040-3052, 2018.
[34] H. Shi, J. Gao, X. Ren, H. Xu, X. Liang, Z. Li, and J. T. Kwok. Sparsebert: Rethinking the importance analysis in self-attention. CoRR, abs/2102. 12871, 2021.
[35] R. Socher, A. Perelygin, J. Wu. J. Chuang, C. D. Manning, A. Y. Ng. and C. Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 conference on empirical methods in natural language processing, pages 1631-1642, 2013,
[36] Y. Tay, M. Dehghani, S. Abnar, Y. Shen, D. Bahri, P. Pham, J. Rao, L. Yang, S. Ruder, and D. Metzler. Long range arena: A benchmark for efficient transformers. CoRR, 1,000000 2020.
12

---

[37] Y. Tay, M. Dehghani. D. Bahri, and D. Metzler. Efficient transformers: A survey, CoRR, abs/2009.06732, 2020.
[38] A. Vaswani, Shazeer, N. Parmar, Uszkoreit, L. Jones, A. N. Gomez, Kaiser, and
I. Polosukhin. Attention is all you need. CoRR, abs/1706.03762, 2017.
[39] A. Wang, A. Singh, J. Michael, F. Hill, O. Levy, and S. R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. CoRR, abs/1 804.07461, 2018.
[40] S. Wang, B. Z. Li, M. Khabsa, H. Fang. and H. Ma. Linformer: Self-attention with linear complexity. CoRR, abs/2006.04768, 2020.
[41] Z. Wang. Sparsert: Accelerating unstructured sparsity on gpus for deep learning inference. arXiv preprint arXiv:2008. 1849, 2020.
[42] A. Williams, N. Nangia, and S. Bowman. broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long Papers), pages 1112-1122, New Orleans, Louisiana, June 2018. Association for Computational Linguistics.
[43] Y. Xiong, Z. Zeng, R. Chakraborty, M. Tan, G. Fung, Y. Li, and V. Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. CoRR, abs/2102.03902, 2021.
[44] C. Yun, Bhojanapalli, A. Rawat, S. J. Reddi, and S. Kumar. Are transformers universal approximators of sequence-to-sequene, functions? CoRR, (bs/1992.10077,, 2019.
[45] C. Yun, Y. Chang, S. Bhojanapalli, A. S. Rawat, S. J. Reddi, and S. Kumar. O(n) connections are expressive enough: Universal approximability of sparse transformers. CoRR, abs/2006.04862, 2020.
[46] M. Zaheer, G. Guruganesh, A. Dubey, J. Ainslie, C. Alberti, S. Ontañón, P. Pham, A. Ravula, Q Wang, L. Yang, and A. Ahmed. Big bird: Transformers for longer sequences. CoRR, 0e//2007.14663 2020.
[47] R. Zellers, A. Holtzman, H. Rashkin, Y. Bisk, A. Farhadi, F. Roesner, and Y. Choi. Defending against neural fake news. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bous Buché-Buc,,a E. Fox, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 32, Curran Associates, Inc., 2019.
[48] P. Zhang, X. J. Yang, B. Xiao, I.. Yuan, L.. Zhang, and Gao. Multi-scale vision longformer: new vision transformer for high-resolution image encoding. In Proceedings of the IEFF/CV International Conference on Computer Vision, pages 2998-3008. 2021.
[49] M. Zhu and Y. Xie. Taming unstructured sparsity on gpus via latency-aware optimization. In 2020 57th ACM/IEEE Design Automation Conference (DAC), pages 1-6, 2020.
[50] Y. Zhu, R. Kiros, R. Zemel, R. Salakhutdinov, R. Urtasun, A. Torralba, and S. Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In Proceedings of the IEEE international conference on computer vision, pages 19-27, 2015.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See
Section 6
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix A.
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes] See Figure in Section and Appendix C.
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [Yes]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any assets either in the supplemental material or as URL? [Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you're rou're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "rAvBA17vVs", "reviewer": "Reviewer_sh4F", "summary": "This paper proposes a Transformer variant where the self-attention matrices are masked using random adjacency matrices generated using a stochastic block model (SBM). The goal here is to reduce the computational cost of producing each N x N self-attention matrix. In the forward pass, the proposed method computes SBM parameters from the query and key embeddings using an MLP and samples a self-attention mask using a fast random graph sampling method. In the backward pass, gradients are passed through this discrete sampling step using a straight-through estimator. The authors show that the proposed SBM-Transformer is a universal approximator in the setting considered by Yun et al. (2020) and Zaheer et al. (2020). Empirically, the authors demonstrate that the SBM-Transformer (1) is able to emulate a standard Transformer with dense attention on a synthetic task and (2) outperforms several efficient Transformer variants on average on the Long Range Arena benchmark.", "questions": "- As discussed in the previous section, does the SBM-Transformer push out the accuracy-cost Pareto frontier vs. existing efficient Transformer architectures?\n- Nondeterministic prediction at inference time is typically considered to be undesirable by practitioners. However, the setup as described requires random sampling under the SBM even during inference. Can this sampling step by replaced with a deterministic operation?\n", "limitations": "The authors highlight the limitation of their work with regard to hardware support for unstructured sparsity in Section 6.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Originality\n\n- As noted in the paper, the use of sparse attention masks to reduce the computational cost of the self-attention operation has previously been explored in several existing papers (e.g., axial attention (Ho et al. 2019), BigBird (Zaheer et al., 2020)). The main methodological novelty here is the use of a stochastic block model to generate a random attention mask in each forward pass. To my knowledge, this approach has not been studied in prior work.\n\n### Quality\n\n- **Theoretical results.** The paper establishes a universal approximation property for the proposed transformer variant (Theorem 1). This result is helpful as a sanity check that the proposed method is a reasonable approach. The overall technical contribution here is relatively minor since the proof of the result is primarily an application of existing machinery from Yun et al. (2020) and Zaheer et al. (2020). A minor weakness of the analysis is that it makes use of a modified version of the SBM-Transformer architecture relative to that used in the empirical evaluation -- in particular, the analysis requires that M_ii = 1, unlike in the version used to obtain the results in Table 1.\n- **Empirical evaluation.** A strength of the paper is the result that the proposed SBM-Transformer architecture outperforms the baseline efficient Transformer architectures on average over the LRA benchmark tasks (Table 1). However, a significant weakness of this comparison is that it does not appropriately account for the accuracy-cost tradeoff corresponding to each of the evaluated methods. For instance, it is not obvious from the paper whether the SBM-Transformer outperforms the given baselines when controlling for a measure of computational cost such as the per-example FLOP count during inference. Such controls are important since the SBM-Transformer introduces additional operators (such as a 2-layer MLP used for computing the SBM parameters), the costs of which will have to be offset by the increased sparsity of the self-attention matrix.\n\n### Clarity\n\n- The presentation was generally clear and easy to follow.\n- The specification of the masked attention computation in Section 3.2 is problematic since simply masking the attention logits to 0 is not equivalent to zeroing-out the attention scores after the softmax operator. \n\n### Significance\n\n- As noted in Section 6, the unstructured sparsity of the resulting attention maps in the SBM-Transformer formulation is not amenable to fast matrix multiplication on current GPU hardware. Thus, as a practical matter, the proposed method requires longer wall clock times for inference, which consequently limits the significance of the work for practitioners. The authors argue that the method (1) \"enables multi-head attention with long sequences\", and that (2) practitioners can tune the computational cost of inference by imposing constraints on the number of edges to sample. I don't find these arguments to be convincing as is -- in particular, claim (1) requires additional justification in light of (a) existing efficient transformer architectures that incur subquadratic computation, and (b) techniques such as memory-efficient attention [1] that require subquadratic memory. As for claim (2), the merit of this additional flexibility (e.g., in terms of accuracy vs. cost) should be evaluated empirically.\n\n[1] Rabe & Staats, 2021. Self-attention Does Not Need O(n^2) Memory. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "ObqUtvGlBU7", "reviewer": "Reviewer_Nfpt", "summary": "The work under revision proposes SBM-Transformer. A novel variant of a transformer architecture adjusts a sparsity of its attention blocks dynamically, based on input sequence data. It is done so forward and backward evaluation cost is linear in the number of edges. Those are adaptively chosen by deploying a generative model for community detection, stochastic block model (SBM), combined with a fast sampling method, fastRG (Alg. 1, developed in [30]). Such an architecture is then flexible and capable of using a mixture of sparsities of attention blocks across the layers yet still capable of learning a full attention model if required. On average, however, it reduces computational complexity significantly, as documented in the well-chosen experimental section on the benchmark Long Range Arena (LRA) dataset. \n\nThe paper builds on previous theoretical expressivity results of sparse attention transformers and applies this result to proposed architecture in two Lemmas, resulting in Theorem 1 with proofs in the Appendix.\n\nThe paper concludes on the limitation of the method. Unfortunately, GPU architectures used to train large transformers are not efficient for sparse tensor operations. Proposed SMB-Transformer thus \" ...often results in longer runtimes compared to dense counterparts even though the memory usage is lower ... \". ", "questions": "To help accessibility of the article (see above) I suggest rewriting/extending the lines between 145 - 164. For instance why bipartite structure? Or what is 2-layer MLPd_h to d_h and why two layers (most likely to retain dimensionality but could you put some motivating lines of thoughts down?)\n\nAre there any numerical or other limitations incurred by the use of SBM? In particular, this generative model assumes graph edges follow (mean parameterized) Poisson distribution. Besides one mentioned (line 182) regarding a degeneration of the mask matrix M, I'd encourage authors to elaborate more on this topic. For instance, is there any limitation in terms of learning \"outliers\" that is small clusters but highly connected clusters/tokens? (SMB is known to work well for the balanced size of clusters, but may struggle for imbalanced ones.)\n\nPlease find more suggestions in the previous block of comments.", "limitations": "A large part of the conclusions is dedicated to a major limitation - despite using fewer memory run-times on GPU clusters are often longer for SBM-Transformer than for dense models. In my opinion, this is unfortunate concerning the practical impact of this work because a majority of large transformer models GPT-3, PaLM, OPT175B, etc. have been trained in more and more accessible High-Performance Computing (HPC) centers, i.e. supercomputers, with massive GPU partitions.\n\nSecondly all results hold \"in expectation\" (in a weak sense). Could authors elaborate more on the numerical stability/deterioration of the algorithm?", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, this is an interesting application of sampling methods (an algorithm fastRG[30]) combined with gradient descent optimization with theoretical efficiency guarantees supported by experiments while not sacrificing the performance. To the reviewer's knowledge, this is a novel approach that I believe is of high interest to the ML community.\n\nShort term, the practical impact/significance of the work is limited due to NOT being well suited for large GPU clusters - the nowadays platform of choice for large transformer models. Also mentioned in a nicely written Related Work, there are several methods in place, e.g., Reformer [18], using hashing similarity instead of SBM, similar in the idea, be it more constrained.\n\nThe paper could also be improved in terms of accessibility and readability, especially for a reader not familiar with SBM and generative modeling. For instance, it could help to elaborate in detail how a stochastic block model version from paper [2] (I believed that one is used), translates to SMB-transformer proposed. Especially in section 4.1.\n\nTheorem 1: The idea of the proof presented in the main body of the paper is still technical and does not spark the intuition of why a “star graph” structure ensures expressibility after rephrasing the idea of the theorem from [38] on line 213. Moreover, the proof in Appendix relies on previous works, mainly by Yun et al.,(2020) [14], and Theorem 1 is proven by showing that SMB-Transformer meets the necessary assumptions of Yun et al.,(2020) [14]. Proof of Lemma 2 shows expressibility for p=3 sparsity patterns (expressibility is not guaranteed for p=2 as noted in line: 215) and thus proves it for higher p’s. However, it rather seems like a direct application of previously mentioned works on this special case of SMB-Transformer. Could you elaborate on what is a novelty here? Otherwise, I suggest being specific and making it an application of Theorem XY from [] and [] ...\n\nline 285 and Table 2. Given the experimental std/errors (in brackets?) presented in the table, I am not sure we can conclude that “… applying a small sparsity regularizer helps in boosting overall sparsity as well as performance …”.  I would claim it is rather insignificant looking at average results in the last column! I agree, however, the computation saving effect of regularization vs. no reg. in Figure 5 still holds. So I suggest considering rephrasing or omitting the claim above perhaps.\n\nMinor comments: \ntypo: typo Table 1 “the the”", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Al9sNHsFmFU", "reviewer": "Reviewer_ukVv", "summary": "In this paper, the authors propose to adaptively learn the sparsity of attention module through modeling the mask matrix as a bipartite graph generated by SBM.\n\nSpecifically, they discuss how they can apply the new attention mechanism to both forward and backward propagation.\n- They use Eqn. (5), (6), (7), and (8) to model the mask matrix in forward propagation, which introduce a trainable cluster embedding matrix $C$ and two MLPs to decide the node membership of tokens in $Q, K$.\n- They \"incorporate a Straight-Through Estimator (STE) to pass the gradient beyond the discrete sampling\", which works as if $M \\odot \\mathbb E M$ is used in forward propagation.\n- They also show the new attention mechanism is a universal approximator, and evaluate their proposed methods on Long-range-arena (LRA) benchmarks.", "questions": "Could the authors show the extra time and space cost of adding the SBM mechanism? How is that compared to another efficient attention mechanism?", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Originality\n\nThe idea to model the sparsity of attention by SBM is new, though formally similar to the classical idea in pruning that directly model the whole mask matrix.\n\n### Quality\n\nPros: \n\n- The author propose a new data-adaptive sparse attention mechanism.\n\n\nCons:\n- The experiments on LRA is not that convincing nowadays given the recent paper \"Efficiently Modeling Long Sequences with Structured State Spaces\" (ICLR 2022), which shows current Transformers are way worse than their state-space-based method. The effectiveness of the experimental results would be in doubt whether in another setting/benchmarks which Transformers are more good at, the conclusion would still hold.\n- The actual runtime of the new method is not reported and compared to exisitng methods.\n\n\n### Clarity\n\nThe clarity of this paper is awesome. There is a corner concern: to my knowledge, \"SBM\" should be short for \"stochastic block model\" rather than stochastic blockmodel. \n\nAlso, the statement of Theorem 1 is not very informative. It would be better if the content in Appendix A that Eqn. (10) can hold with $O(n)$ connections can be moved to the main paper.\n \n### Significance\n\nThe problem studied is important: how to sparsify and therefore accelerate attention. However, the GPU-unfriendly sparse matrix multiplication makes the significance limited. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept