  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 .
