Transformers meet Stochastic Blockmodels: Attention with Data-Adaptive Sparsity and Cost
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 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 problems byendowing each attention head with a mixed-membership Stochastic Blockmodel(SBM).
Then, each attention head data-adaptively samples a bipartite graph, theadjacency of which is used as an attention mask for each input.
During backpropagation, a straight-through estimator is used to flow gradients beyond the discretesampling step and adjust the probabilities of sampled edges based on the predictiveloss.
The forward and backward cost are thus linear to the number of edges, whicheach attention head can also choose flexibly based on the input.
By assessing thedistribution of graphs, we theoretically show that SBM-Transformer is a universalapproximator for arbitrary sequence-to-sequence functions in expectation.
Empirical evaluations under the Long Range Arena benchmark demonstrate that ourmodel outperforms previous efficient variants as well as the original Transformerwith full attention.
Introduction
The Transformer [34] architecture has been the go-to method for encoding sequential data, due to itssuperior performance in various tasks such as machine translation [26], image classification [12] andprotein language modeling [29].
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 [14].
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 [15] or high-resolution image processing [39].
Many works have thusfocused on developing more efficient alternatives by exploiting fixed or learnable attention sparsitypatterns [7, 38, 18, 11], low-rank approximations [35, 36], or kernelized attention modules [17, 8].
Even though the efficient alternatives hold theoretical expressibility guarantees [37], 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 unanimouslysparsify all the attention modules across layers, without considering the tasks at hand.
Such strategy
[Figure]
Linear [Figure] 𝑸 𝑲 𝑽 imposes inductive bias too strongly and often leads to sub-optimal cost vs. performance trade-offs indownstream tasks [25].
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 [9].
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 Blockmodel (SBM) [2], each attentionhead samples a bipartite graph connecting queries to keys.
Then, the adjacency of the sampledgraph is used as an attention mask so that only attention scores corresponding to sampled edgesare computed.
The overall computational cost is linear in the number of edges, which can rangefrom linear to quadratic in sequence length depending on the data and task under concern.
Eachattention head is equipped with its own underlying SBM, enabling the model to diversify the attentionsparsity across heads and layers.
By incorporating a straight-through estimator [4] in the discretegraph-sampling step, SBM-Transformer enjoys end-to-end differentiability and can find the properattention sparsity based solely upon minimizing the predictive loss.
The model can also easily befurther 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 where full attention is required to achieve 100% accuracy.
• Extensive evaluations on the LRA benchmark 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 [33].
The Linear Transformer [17]achieves linear complexity by replacing the softmax with a low-rank kernelized function.
Linformer [35] and Nyströmformer [36] use a similar approach by low-rank approximating the attentionscore matrix.
Performer [8] uses positive orthogonal random features to approximate the softmax kernel.
Reformer [18] gathers similar tokens together through locality-sensitive hashing (LSH) andperforms 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 [10] proposed replacing the usual softmax activation with α -entmax,in which the α parameter can be differentiably trained to adjust the activation between softmaxand sparsemax activation [23].
SparseBERT [31] 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 Blockmodels.
The Stochastic Blockmodel (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 [20] 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 [16]for multi-graphs and hierarchical SBM [27] for multiplex-graphs.
Further details can be found in arecent survey [13].
Preliminaries: Sparse Transformers
We first introduce the full attention mechanism used in the original Transformer [34] as well asmasked attention which will serve as a backbone of our approach.
3.1 Full Attention
In vanilla Transformer [34], 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 of
𝑸 𝑲 𝑽 ෡𝑸 ෡𝑲 𝑪 ෡𝑺
Attention Mask
GAT [Equation] 𝜕ℒ/𝜕𝑿
Forward
Backward
I n p u t s
Clusters
Query/KeyNodes
SBM +
STE
Output 𝑸,𝑲 [Equation] where ⊙ indicates entry-wise multiplication.
Note the normalizing term in masked-softmax σ M ( · )only sums up unmasked terms, ensuring that each ( i, j ) -th attention score survives as nonzero if andonly if M ij = 1 .
Most sparsity-based efficient Transformers fall under this formulation, while usingdifferent methods to either manually fix or learn the mask M .
For instance, local attention [7, 3, 38]with a sliding window sets M ij = 1 if | i − j | < c for some context window size c while Reformer [18]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 Blockmodel
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 blockmodels 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 a [Figure]
(cid:2417)(cid:4670)(cid:2169)(cid:4671)
[Figure] 2-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 connectingquery Q i to key K j equals p ( i, j ) =
ˆ Q i ˆ S ˆ K Tj .
Formally, the graph sampling procedure can bewritten as
[Equation]
[Equation] [Equation] [Equation] For the last sampling step, we incorporate a fast random graph sampling algorithm fastRG (Alg.
1,[30]) 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 [18].
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, a naive backpropagationcan not 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 basedon the loss. The gradients provide information on whether each sampled edge was useful or not forprediction.
Algorithm
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 ∂ L [Equation] if M ij = otherwise where A := M ⊙ QK T √ d h (9) 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 theprediction 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
Following the theoretical analysis of BigBird [38], we demonstrate that the mixed-membership SBMwith adaptability and stochasticity provides the same level of expressibility as full attention.
Forbrevity, we briefly discuss the main idea of the proof, and defer the details to the appendix.
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]
[Figure]
Theorem 1 from [38] states that sparse attention enjoys expressibility if the star graph with one globalrelay token is contained within the sparsity graph. In our setting, the star graph can be viewed as arank-2 graph structure where the global relay token is in its own cluster, and all other tokens belongto another.
As long as each attention head uses more than 2 clusters, there exists an SBM which cansample the star graph in expectation, hence our model achieves universal approximability.
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) [32], a benchmark widely used to assess the capacity of efficient Transformersin learning long-range contexts across different modalities.
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 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.
[Table]
⇒
[Table]
Methods.
For this task, we compare SBM-Transformer with k = 128 clusters against variousefficient Transformers: Linear Transformer [17], Linformer [35], Reformer [18], Performer [8], and
Nyströmformer [36].
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 attentionare unable to achieve the same level of accuracy, due to limited expressibility under the constrainedsetting.
Though SBM-Transformer converges more slowly compared to full-attention, it demonstratesthe ability 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 in modeling longrange dependencies, we test SBM-Transformer against previous work on the LRA benchmark. We
[Table] also test how the performance is affected with respect to changing number of clusters of the SBMsand applying a sparsity-based regularizer.
Dataset.
LRA [32] consists of five different testbeds with varying modalities: L IST O PS [24] is a10-way classification task to map a sequence to its corresponding solution. Each sequence consists ofsingle-digit numbers and 4 different set operations.
T EXT [22] is a binary classification task wherebyte-level IMDB movie reviews must be classified into one of two sentiments: positive or negative.
R ETRIEVAL [28] is also a char-level binary classification task. Unlike T EXT , two sequences from
ACL Anthology papers are given as input, and the model must predict whether there exists a citationlink between them.
I MAGE [19] is a 10-way classification task mapping flattened pixel-sequencesfrom CIFAR-10 to its class.
P ATHFINDER [21] provides flattened pixel-sequences from an image andthe model must decide whether two circles are connected by a dashed line.
For this benchmark, weuse the PyTorch implementation of LRA provided by the authors of Nyströmformer [36] and adhereto the same train-test splits.
Performance in all five tasks is measured using classification accuracy.
Methods.
We compare SBM-Transformer against the same baseline methods as with the synthetictask above.
For fair comparison, we set all Transformer models to use the default setting suggestedin previous work [36], which fixes 2 layers, 2 attention heads, and 64 embedding dimensions.
Theoutput token representations are mean-pooled to obtain the sequence representation for all tasks. For SBM-Transformer, we use k = 128 clusters in all attention heads.
More details on the architecturesetups can be found in the appendix.
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 tradeoff betweencomputational cost and performance, we also test the model with the same hyperparameter settingunder a sparsity-based regularizer that discourages excessive use of query-key edges.
We penalizeeach sampled edge by adding to the predictive loss a weighted regularization term λ L s , where L sdenotes the average density of the sampled mask across all attention heads in the model.
Table 2shows the performance of SBM-Transformer across varying regularization weights λ .
Surprisingly,the model retains competitive performance under the constrained setting compared to the modeltrained without any regularization, while reducing the average mask density significantly.
Themodel even outperforms the unregularized model when trained under a small regularization weightλ = 10 − 4 , indicating that the signal from λ L s does not severely interfere with that from predictionloss.
It may also be the case similar local optima are shared across sparse attention regions withvarying density in the loss landscape, and the regularization term is able to drive the model towardsfinding optimal attention scores with smaller density.
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 sparsityacross layers for better performance.
Recall that this allows models to gather information in differentlevels, as seen in pretrained BERT where lower layers focus on the overall content via dense attention
[Table]
[Figure]
while upper layers gather syntactic information with tree-like patterns [9].
For each of the five tasks,we pick two highest-performing models (one for unregularized and another for regularized) formeasurement.
Figure 5 shows the average layer-wise mask densities of unregularized and regularized
SBM-Transformers across different tasks.
We find that under no regularization, the two layers candiffer by more than 10% in tasks such as L IST O PS and I MAGE .
This may be due to the hierarchicaland compositional structure of the two tasks.
We also find that the variation is relatively low in
T EXT with densities around 25%, indicating that the task requires broad attention overall. Lastly,the standard deviation is extremely large in upper layers for P ATHFINDER , showing that it samples awide variety of masks depending on the input.
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 beforehand.
Theoretically, we show that our model enjoys the same expressibility as theoriginal Transformer due to the flexibility and stochasticity of the latent SBM.
Empirical experimentson LRA show that our model outperforms previous state-of-the-art efficient Transformers.
Nonetheless, there are limitations from an engineering perspective as operations with sparse tensorsare not GPU-friendly.
It often results in longer runtimes compared to dense counterparts eventhough the memory usage is lower.
Many sparsity-based attention mechanisms have instead usedblock-sparse attention that is much more amenable for GPU computation [38, 7, 3].
We still believethis work is valuable as it enables multi-head attention with long sequences, and practitioners caneasily impose constraints on the number of edges to sample, based on the computational resourceavailable in their environment.
We hope to see more GPU-friendly tensor operations optimized forfine-grained sparsity in the future, at which point the value of this work will increase even further.
For future work, we hope to extend this work further towards a higher-level Neural Architecture
Search (NAS) in the context of Transformers where the model not only flexibly tunes the attentionsparsity, but also adjusts the number of layers and attention heads during training to autonomouslyoptimize the tradeoff between model size and predictive performance.
As we propose a foundationalreplacement for the scaled dot-product attention module in the Transformer architecture, we do notexpect any immediate negative societal impact due to this work.