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 this https URL .

本研究提出了 SBM-Transformer 模型，利用基于混合成员资格随机块模型（SBM）的稀疏注意力来解决二次成本问题，并证明了 SBM-Transformer 是任意序列到序列函数的通用逼近器，同时在 LRA 和 GLUE 基准测试中表现出优异的性能。