Transformers meet Stochastic Block Models: Attention with Data-Adaptive Sparsity and Cost s AffiliationAddress
Abstract
To overcome the quadratic cost of self-attention, recent works have proposedvarious 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 asparse variant of softmax such as α -entmax.
Unfortunately, the first group lacksadaptability to data while the second still requires quadratic cost in training.
Inthis work, we propose SBM-Transformer, a model that resolves both problemsby 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.
Duringbackpropagation, a straight-through estimator is used to flow gradients beyondthe discrete sampling step and adjust the probabilities of sampled edges basedon the predictive loss.
The forward and backward cost are thus linear to thenumber of edges, which each attention head can also choose flexibly based onthe input.
By assessing the distribution of graphs, we theoretically show that SBM-Transformer is a universal approximator for arbitrary sequence-to-sequencefunctions in expectation.
Empirical evaluations under the Long Range Arenabenchmark demonstrate that our model outperforms previous efficient variants aswell as the original Transformer with full attention.
Introduction
The Transformer [38] architecture has been the go-to method for encoding sequential data, due to itssuperior performance in various tasks such as machine translation [28], image classification [14] andprotein language modeling [32].
Its key strength stems from the multi-head attention module, wherea so-called attention score matrix computes how contextually important one token is to another for allpossible token pairs.
Each Transformer layer simultaneously pools the token representations basedon the attention scores, eventually returning contextualized features without sequentially traversingthrough the input sequence as its recurrent neural network-based predecessors [16].
A well-known drawback of the original Transformer is its high computational cost in time and memorythat increases quadratically with sequence length.
This is due to the full pairwise computation ofattention scores, which prohibits applying it in tasks involving long-range dependencies such asdocument summarization [17] or high-resolution image processing [48].
Many works have thusfocused on developing more efficient alternatives by exploiting fixed or learnable attention sparsitypatterns [8, 46, 20, 12], low-rank approximations [40, 43], or kernelized attention modules [19, 9].
Even though the efficient alternatives hold theoretical expressibility guarantees [45], they are farfrom sufficient, still failing to convince practitioners to replace the original Transformer.
We believethis is mostly due to their lack of adaptability.
They apply the same modifications to unanimously sparsify all the attention modules across layers, without considering the tasks at hand
Such strategy
[Figure]
𝑸 𝑲 𝑽 imposes inductive bias too strongly and often leads to sub-optimal cost vs. performance trade-offs indownstream tasks [27].
In this work, we argue that to retain the utmost potential of the 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 amixture of dense and sparse attention layers.
For example, a qualitative analysis on pretrained BERTshowed that lower layers exhibit broad dense attention while upper layers perform focused sparseattention [10].
In the case of GPT-3 [6], the Transformer blocks are manually arranged to alternatebetween dense and sparse attention.
To contribute to the efficient Transformers lineage, we propose SBM-Transformer, capable ofadjusting its attention sparsity data-adaptively based without fully computing the attention scorematrix (Figure 1).
Leveraging a mixed-membership Stochastic Block Model (SBM) [2], eachattention head samples a bipartite graph connecting queries to keys.
Then, the adjacency of thesampled graph is used as an attention mask so that only attention scores corresponding to samplededges are computed.
The overall computational cost is linear in the number of edges, which canrange 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 theattention sparsity across heads and layers.
By incorporating a straight-through estimator [4] in thediscrete graph-sampling step, SBM-Transformer enjoys end-to-end differentiability and can find theproper attention sparsity based solely upon minimizing the predictive loss.
The model can also easilybe further regularized by penalizing the number of sampled edges, which results in a lighter modelusing less computational resources during inference.
To the best of our knowledge, our method is thefirst Transformer architecture that can data-adaptively choose between linear to full attention withrespective computational costs.
To summarize, our main contributions are as follows:
• We present SBM-Transformer, a 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 SBM-Transformerretains universal approximability, and also stress-test the model under a synthetic task wherefull attention is required to achieve 100% accuracy.
• Evaluations on LRA and GLUE benchmarks show that SBM-Transformer outperforms previous efficient Transformer models as well as the vanilla Transformer with dense attention.
Related Work
In this section we discuss previous efficient Transformer variants and several works similar to ourswith 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 multihead attention with different approaches.
While we provide only a handful of representative approaches, a 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.
Linformer [40] and Nyströmformer [43] use a similar approach by low-rank approximating the attentionscore matrix.
Performer [9] uses positive orthogonal random features to approximate the softmaxkernel.
Reformer [20] gathers similar tokens together through locality-sensitive hashing (LSH) and performs attention amongst tokens within the same bucket.
Of all methods above, our method ismost similar to Reformer, in the sense that we adaptively assign queries and keys into clusters andform a low-rank sparse attention pattern.
However, our method performs soft-clustering with muchless structural constraints, allowing each attention head to represent a wider variety of dependencystructure and to adjust its sparsity towards full attention if needed.
Adaptive Sparsity.
With respect to flexible training between sparse and dense attention, there existsome works that parameterize how sparse the attention pattern should be based on the input.
The
Adaptive Sparse Transformer [11] proposed replacing the usual softmax activation with α -entmax,in which the α parameter can be differentiably trained to adjust the activation between softmaxand sparsemax activation [25].
SparseBERT [34] uses a differentiable masking technique whereeach attention mask is sampled from a Gumbel-sigmoid distribution using data-independent maskprobability parameters.
While these methods possess the flexibility to adjust between sparse andfull attention based on data, they still require full computation of the attention score matrix beforesparsification, 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 be able to adaptively tune its attention sparsitybetween sparse to full attention without requiring the explicit computation of the attention scorematrix, thereby avoiding quadratic cost when possible.
Stochastic Block Models.
The Stochastic Block Model (SBM) is a generative model that encodesthe latent structure of graphs by grouping nodes into clusters.
By modeling the cluster-membershipof each node as well as inter-cluster relationships, SBMs can represent a wide variety of graphstructures, which is a feature especially useful for generating new graphs or predicting missingedges in noisy data [1].
The standard SBM assigns each node to a single cluster, and the probabilityof an edge between two nodes strictly depends on the corresponding clusters.
Several structuralextensions include overlapping SBM [22] and mixed-membership SBM [2], which allow each node tobe assigned to multiple clusters.
The underlying SBM used by our framework mostly resembles thesetwo variants, while the edge probability is modeled by a nonlinear function of two node embeddingsrather than a bilinear one.
There exist many other extensions including degree-corrected SBM [18]for multi-graphs and hierarchical SBM [29] for multiplex-graphs.
Further details can be found in arecent survey [15].
Preliminaries: Sparse Transformers
We first introduce the full attention mechanism used in the original Transformer [38] as well asmasked attention which will serve as a backbone of our approach.
3.1 Full Attention
In vanilla Transformer [38], each attention head takes a sequence of token features as input X ∈ R n × dwhere n is the sequence length and d the embedding dimension.
Weight parameters W Q , W K ∈
111 R d × d h and W V ∈ R
d × d h with head-dimension d h 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-productof queries and keys followed by row-wise softmax activation σ ( · ) .
Note that explicit computation ofthis matrix is the main bottleneck of full attention, incurring O ( n 2 ) asymptotic cost in both time andmemory.
The value features V are then pooled based on the attention scores, returning the outputtoken representations.
Altogether, the operation performed by each attention head can be written as
[Equation]
3.2 Masked Attention
One way to remove the quadratic bottleneck from the attention score matrix is to apply a binary mask
M ∈ { 0 , 1 } n × n and compute the scaled dot-products Q
i K j / √ d h only if M ij = 1 .
In presence ofan attention mask, the operation is modified to [Equation]
𝑸 𝑲 𝑽 ෡𝑸 ෡𝑲 𝑪 ෡𝑺
Attention Mask
GAT 𝑿 𝜕ℒ/𝜕𝑿
Forward
Backward
I n p u t s
Clusters
Query/KeyNodes
SBM +
STE
Output 𝑸,𝑲 where ⊙ indicates entry-wise multiplication.
Note the masked-softmax σ M ( · ) operator onlycomputes unmasked terms, ensuring that each ( i, j ) -th attention score survives as nonzero if and onlyif M ij = 1 . This is thus equivalent to filling in the ( i, j ) -th attention score with −∞ if M ij = 0 ,then applying the standard softmax operator.
Most sparsity-based efficient Transformers fall underthis formulation, while using different methods to either manually fix or learn the mask M .
Forinstance, local attention [8, 3, 46] with a sliding window sets M ij | i − j | c window size c while Reformer [20] sets M ij = 1 if Q i and K j are hashed into the same bucket.
Our Method: SBM-Transformer
Here we discuss the details of SBM-Transformer (Figure 2).
We first illustrate the forward step of ourattention module and how the underlying SBM [2] of each head, from which we sample our attentionmasks, is parameterized by the input tensors.
We then discuss how the model enables end-to-enddifferentiability 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 thatconnects queries to keys, and let each attention head sample an adjacency matrix that best representsthe contextual dependencies amongst input tokens.
In order to efficiently sample adjacency matriceswhile avoiding the quadratic cost, the distribution of graphs must first be parameterized with asub-quadratic number of latent variables.
Stochastic Block Models fit perfectly for our purpose asit models graphs that are low-rank structured with k latent clusters, allowing full parameterizationwith O ( nk ) memory.
More concretely, the SBM distribution is defined by two nonnegative nodeto-cluster memberships Y , Z ∈ R n × k
+
and a so-called block matrix B ∈ R k × k + that stores theinter-cluster connection probabilities.
The probability of node i being connected to node j iscomputed as p ( i, j ) =
Y i
BZ Tj .
Equivalently, the expectation of the adjacency matrix sampled from
A ∼ SBM ( Y , B , Z ) can be written as E
[ A ] =
Y BZ T .
For proper parameterization of the SBM, we must infer the nonnegative node-memberships and blockmatrix from the queries and keys.
To do so, we equip each attention head a 2-layer MLP d h → d h with
ReLU activation, and a set of k trainable cluster-embeddings C ∈ R k × d h .
First, our model computesthe block matrix ˆ S ∈ R k × k + by taking dot products amongst cluster-embeddings C followed by a2-dimensional softmax activation.
The node embeddings are obtained by processing each query andkey through the MLP d h → d h , mapping token representations into the node representation space.
Thememberships of query and key nodes, which we denote by ˆ Q and ˆ K , are then inferred by takingdot products of node and cluster embeddings, followed by a sigmoid function.
The block matrixˆ S , query node-memberships ˆ Q , and key node-memberships ˆ K altogether provide a well-definedparameterization for the SBM.
Thus, a bipartite graph adjacency M ∈ { 0 , 1 } n × m can be sampledfrom M ∼ SBM ( ˆ Q , ˆ S , ˆ K ) with expectation E [ M ] = ˆ Q ˆ S ˆ K T : the probability of connecting [Figure]
(cid:2417)(cid:4670)(cid:2169)(cid:4671) [Figure] query Q i to key K j equals p ( i, j ) =
ˆ Q i ˆ S ˆ K Tj .
Formally, the sampling procedure can be written as
[Equation]
For the last sampling step, we incorporate a fast random graph sampling algorithm fastRG (Alg.
1,[33]) that can sample graphs from a SBM in time and memory asymptotically linear in the numberof edges.
One advantage of fastRG is that each edge can be sampled in parallel, allowing highefficiency with the help of multiprocessing.
A more significant feature of the method is that thenumber of edges, which determines the overall cost, is sampled from a Poisson distribution
withinput-dependent mean (Line 4).
Thus, the model can dynamically adjust its computational costbetween linear and quadratic in sequence length based on the data.
Figure 3 shows example placements of nodes and clusters on the d h -dimensional space to show howthe sparse structure is determined.
If all nodes and clusters are gathered closely, then all entries in ˆ Qand ˆ K become close to 1, resulting in p ( i, j )
≈
for all i, j and hence a dense M .
If clusters arewell-separated but each surrounded by some set of nodes, ˆ S becomes close to diagonal while eachrow in ˆ Q and ˆ K is close to a one-hot vector indicating the cluster nearby.
Such setting leads to ablock diagonal mask similar to LSH bucketing of Reformer [20].
Lastly, if all clusters are far apartfrom the nodes, both ˆ Q and ˆ K approximately equal zero, zeroing out all the edge probabilities.
4.2 Backward Step with Straight-Through Estimator
The graph sampling procedure is naturally a discrete operation.
Thus, naive backpropagation cannot learn the proper parameterization for the SBM that minimizes the predictive loss.
To copewith this non-differentiability, we incorporate a Straight-Through Estimator (STE)
[4] to pass thegradient beyond the discrete sampling step.
The STE enables providing the gradient
∂ L /∂ M ij tothe probability for each sampled edge ( i, j ) (Eqn. 9).
It works as if we had used a continuous mask
M ⊙ E
[ M ] that stores the probability of each sampled edge instead of the binary mask M duringforward propagation.
This way, the probabilities of sampled edges can be learned end-to-end: thegradients provide information on whether each sampled edge was useful or not for prediction. [Equation] otherwise where A :=
[Equation] ⊙ QK T [Equation] [Equation]
Random Edge Exploration.
While this approach enables backpropagation in the same O ( m ) costas in the forward step, this comes at the expense of not being able to propagate information throughedges not sampled.
This can be problematic when an edge probability accidentally collapses tozero, after which the edge becomes unlikely to ever be sampled even when it may be useful for the
Input : Y ∈ R n × k + , B ∈ R
k
×
k
+ , Z ∈ R n × k +
Output : M ∈ { 0 , 1 } n × n with E [ M ] =
Y BZ T 1 Compute diagonal matrices D Y =
( diag ( 1 Y ))
− 1 and D Z = ( diag ( 1 Z ))
−2 Column-normalize Y =
Y D − 1 Y and Z = ZD − 1 Z 3 Compute B = D Y BD Z 4 Sample number of edges m ∼ Poisson ( 1 B 1 T )
Initialize M =6 for i = 1 : m do
Sample ( U, V ) from { 1 , . . .
, k } × { 1 , . . .
, k } with Pr ( U = u, V = v ) ∝ B uv
Sample source I from { 1 , . .
. , n } with Pr ( I = i )
=
Y iU .
Sample destination J from { 1 , . . . , n } with Pr ( J = j ) =
Z jV
Set M IJ = 1 .
11 end prediction task at hand.
Therefore, we add a small perturbation δ > 0 to each edge probability p ij ,allowing the model to explore new edges and resuscitate their sampling probabilities if necessary.
We find that a δ as small as 0 .
01 significantly helps in practice, and thus use this exploration schemeduring training for our experiments.
Wouldn’t the model always prefer full attention?
Note that the gradient ∂ L /∂p ij can be positive,which suppresses the probability of edge ( i, j ) .
At first, it may seem counter-intuitive why themodel would ever limit itself to using fewer edges during training.
One explanation is that maskedattention provides an easy way to reduce attention scores under finite head dimensions.
Under fullattention, it is known that the representational space of attention score matrices is limited by the headdimension and softmax activation [5].
This limitation inevitably introduces unwanted noise in theattention scores especially when working with long sequences.
In SBM-Transformer, however, thestructural sparsity in masked attention introduces another dimension that induces a larger space ofrow-stochastic matrices (full attention is a special case of masked attention where M ij = 1 for alli, j ).
Therefore, it is reasonable that the model may encourage sparsity to leverage the additionalexpressiveness assuming the loss landscape has local optima within the sparse attention regime.
Ourexperiments on the LRA benchmark show that this is indeed the case, as our SBM-Transformerconverges to an average attention sparsity of 20% to 30% while outperforming Transformer with fullattention.
We also show in the experiment that we can easily incorporate additional regularizationthat further encourages the sparse mask.
4.3 SBM-Transformer is a Universal Approximator
Leveraging previous work on the theoretical expressiveness of sparse attention [45, 46], we show that SBM-Transformer with a small modification 1 retains the same level of expressibility as full attention.
Specifically, we show that the low-rank structure of the underlying SBMs does not degrade theexpressive power of Transformer, and that SBM-Transformer can universally approximate arbitraryfunctions with O ( n ) connections.
For brevity, we provide a rough overview of the proof and deferfurther details to Appendix A.
Theorem 1.
Let f ∈ F be class of continuous sequence-to-sequence functions.
T h,r,mSBM denote theclass of SBM-Transformers with h attention heads, m head dimension, and r dimensions in hiddenlayers.
Then for any ϵ > 0 and 1 ≤ p < ∞ , there exists a function g ∈ T h,m,rSBM such that
[Equation]
According to the main theorem of Yun et al.
[44], SBM-Transformer achieves universalapproximability if 1) each node attends to itself, 2) the aggregation of all attention patterns containsa Hamiltonian path, and 3) there exists a path between all node pairs.
While the first condition is
[Figure]
trivially true due to our modification, the other two conditions require careful choice of three SBMs. Here we first parameterize one SBM to hard-assign tokens into k equally-sized clusters, inducing ablock-diagonal attention pattern. The other two SBMs are parameterized such that the two graphstogether form a star graph with k global relay tokens.
Combining the three attention patterns lead to aparameterization of SBM-Transformer that satisfies all three conditions, hence proving the theorem.
Experiments
For empirical evaluations, we first use a synthetic task to show that our model is flexible enough tolearn towards full attention when needed in contrast to previous works.
We then experiment on Long
Range Arena (LRA) [36], a benchmark widely used to assess the capacity of efficient Transformersin learning long-range contexts across different modalities.
Lastly, we show results on the GLUEbenchmark [39] to assess the performance of SBM-Transformer in a downstream NLP setting.
Allexperiments 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 sequenceconsists of N integers, each of which is uniformly sampled from { 1 , 2 , . . .
, N } .
We use N = 256 inour setup.
The prediction target is a sequence of equal length, where each token is labeled 1 if thereexists a duplicate somewhere within the sequence, and 0 otherwise.
Below is a simple example with
N = 8 that illustrates the task.
We measure the performance of models via binary cross-entropy loss.
[Equation]
Methods.
For this task, we compare SBM-Transformer with k = 128 clusters against variousefficient Transformers: Linear Transformer [19], Linformer [40], Reformer [20], Performer [9], and
Nyströmformer [43].
Across all methods, we use a single-layer and single-head architecture with 32hidden dimensions.
Note that due to this constrained setting, the sole head must perform full attentionto compare each token to all the others and attain 100% accuracy.
All models are trained for 2000epochs where a new batch of sequences is sampled for every epoch.
We use a batch size of 256 andlearning rate of 1e-3.
Results.
Figure 4 shows the training loss curves of each baseline method as well as SBMTransformer.
Full attention quickly converges to 100% accuracy as expected since it computesall possible pairwise interactions by default.
Other models that apply low-rank or kernelized attentionfail to achieve the same level of accuracy, due to limited expressibility under the constrained setting.
Though SBM-Transformer converges more slowly compared to full-attention, it demonstrates theability to drive itself towards full-attention, eventually attaining zero loss.
5.2 Long Range Arena (LRA)
To demonstrate that the flexible inductive bias of SBM-Transformer is effective for modeling longrange dependencies, we test SBM-Transformer against previous work on the LRA benchmark.
Wealso 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: L IST O PS [26] is a10-way classification task to map a sequence of single-digit numbers and 4 different set operations,to its corresponding solution.
T EXT [24] is a binary classification task where byte-level IMDBmovie reviews must be classified into one of positive or negative sentiments. R ETRIEVAL [30] [Table]
[Table]
is also a char-level binary classification task, where two sequences from ACL Anthology papersare given as input, and the model must predict whether there exists a citation link between them.
I MAGE [21] is a 10-way classification task mapping flattened pixel-sequences from CIFAR-10 toits class.
P ATHFINDER [23] provides flattened pixel-sequences from an image and the model mustdecide whether two circles are connected by a dashed line.
For this benchmark, we use the PyTorchimplementation of LRA provided by the authors of Nyströmformer [43] and adhere to the sametrain-test splits.
Performance in all five tasks is measured using classification accuracy.
Methods.
We compare SBM-Transformer against the same baselines as with the synthetic taskabove.
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, weuse k = 128 clusters. The output token representations are mean-pooled to obtain the sequencerepresentation for all tasks.
More details on the architecture setups can be found in Appendix C.
Results.
Table 1 shows the test accuracies of each method.
Our SBM-Transformer achieves thebest overall performance, ranking first in two tasks, and second in one other.
SBM-Transformer alsooutperforms full attention in all five tasks while computing 30% or less attention scores on average,which supports our claim that masked attention with partial attention score computations can bepreferred over full attention depending on the task.
With respect to the attention mask structure, wefind that flexibility of SBM is indeed beneficial, as Reformer struggles in L IST O PS , most likely dueto the inability of block-diagonal masks to model hierarchical contexts.
Mask Density Regularization.
To test if the model can effectively optimize the cost vs. accuracytradeoff, we also test the model under a sparsity-based regularizer that discourages excessive useof query-key edges.
We penalize each sampled edge by adding to the predictive loss a weightedregularization term λ L s , where L s denotes the average mask density across all attention heads.
Table 2 shows the performance of SBM-Transformer across varying regularization weights.
Understrong regularization, the model surprisingly retains competitive performance while significantlyreducing the average mask density.
This indicates that similar local optima are shared across regimeswith varying attention density in the loss landscape, and the regularization term is able to drive themodel towards finding optimal attention scores with smaller density. Efficiency. Furthermore, we compare computational costs during inference by measuring FLOPcount and peak memory usage.
For SBM-Transformer, we test the model trained under λ = 10 − 1 . Due to lack of support for sparse tensor operations in existing FLOP-counters, we measure FLOPsby manually enumerating through each tensor operation. Table 3 shows that SBM-Transformer iscomparably efficient across all tasks except for T EXT , where SBM-Transformer showed the largestaverage mask density. Note that while the cost of other baselines are fixed after initialization, the costof SBM-Transformer is data-adaptive and can vary input-by-input.
Further analysis and qualitativeexamples demonstrating the input-dependent attention mask densities can be found in Appendix C. [Table] 5.3 General Language Understanding Evaluation (GLUE) To check whether its strong performance demonstrated in LRA extends to the downstream NLPsetting as well, we 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 themodel must predict their positive or negative sentiments. For QQP [7], the task is to determinewhether one question is a paraphrase of the other given a pair of questions. MNLI [42] consists ofsentence pairs, each with a target label indicating whether the two sentences are connected throughentailment, contradiction, or neither.
QNLI [31] consists of sentence-question pairs and the task isto determine whether the sentence contains an answer to the question. Each task is formulated assequence classification, and we measure performance by F1 score on the respective validation sets.
Methods.
Following previous work [43], we arrange a small variant of BERT [13] with 4 layers, 8attention heads, and 512 embedding dimensions. We replace full attention with each attention moduleused in previous experiments.
We use the same hyperparameter settings as in LRA, and apply nosparsity regularization (i.e. λ = 0 ) for SBM-Transformer. We first pretrain each model under themasked language modeling objective for 50 epochs on a 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 the appendix.
[Table] Results. Table 4 reports the F1 scores ofeach method on different NLP tasks.
SBMTransformer performs competitively against fullattention overall, and outperforms all baselinesin SST-2 and QQP. We also find that thefine-tuned SBM-Transformer models use 13.5%dense attention masks on average across alltasks, showing that the model can encode useful information from input sentences effectivelyunder highly sparse attention.
Conclusion
We propose SBM-Transformer, an efficient Transformer that can data-adaptively choose its attentionsparsity between sparse and full attention without the need to explicitly compute the full attentionscore matrix.
Theoretically, we show that our model enjoys the same expressibility as the original Transformer due to the flexibility of the latent SBM.
Empirical experiments on LRA and GLUE showthat our model outperforms previous state-of-the-art efficient Transformers.
Nonetheless, there are limitations due to sparse tensor operations being less optimized on GPU kernels.
In the LRA experiments, we found that SBM-Transformer can result in longer runtimes compared todense counterparts while its memory usage is much lower.
While previous sparsity-based attentionmechanisms with block-sparse attention are much more amenable for GPU computation [46, 8, 3],our work requires an architecture with better workload balancing and acceleration under unstructuredsparsity, for which there is 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 on the input. The cost being dependenton the number of edges also allows practitioners to easily impose constraints based on the availablecomputational resources.
We hope to see more GPU-friendly tensor operations optimized for finegrained sparsity in the future, at which point the value of this work will increase even further.
As wepropose a foundational replacement for the scaled dot-product attention module in the Transformerarchitecture, we do not expect any immediate negative societal impact due to this work.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions 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] SeeSection 6.
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
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 supplementary.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 as a 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 Table 2 and Figure ??
in Section 5.
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)?
[Yes]
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 new assets either in the supplemental material or as a URL?
[Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you’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]
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 [Table]
Board (IRB) approvals, if applicable?
[N/A]
[Table]
c)
Did
you include the estimated hourly
wage
paid to participants and the total
amount
[Table]
spent
on participant compensation?
[N/A
[Table]