# Transformers meet Stochastic Block Models: Attention with Data-Adaptive Sparsity and Cost

arXiv: https://arxiv.org/abs/2210.15541 · HTML: https://mlsanity.com/paper/2210.15541/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 56 references that are in the map.
- First posted on arXiv 27 Oct 2022.

## Builds on (most cited first, 30 of 39)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books](https://mlsanity.com/paper/1506.06724/) (2015)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [Defending Against Neural Fake News](https://mlsanity.com/paper/1905.12616/) (2019)
- [Nystr"omformer: A Nystr"om-Based Algorithm for Approximating Self-Attention](https://mlsanity.com/paper/2102.03902/) (2021)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Efficient Attentions for Long Document Summarization](https://mlsanity.com/paper/2104.02112/) (2021)
- [Community Detection and Stochastic Block Models](https://mlsanity.com/paper/1703.10146/) (2017)
- [Multi-Scale Vision Longformer: A New Vision Transformer for High-Resolution Image Encoding](https://mlsanity.com/paper/2103.15358/) (2021)
- [Scaling Neural Machine Translation](https://mlsanity.com/paper/1806.00187/) (2018)
- [Adaptively Sparse Transformers](https://mlsanity.com/paper/1909.00015/) (2019)
- [ListOps: A Diagnostic Dataset for Latent Tree Learning](https://mlsanity.com/paper/1804.06028/) (2018)

## Built on by (most cited first, 4 of 4)

- [Curve Your Attention: Mixed-Curvature Transformers for Graph Representation Learning](https://mlsanity.com/paper/2309.04082/) (2023)
- [Neural-prior stochastic block model](https://mlsanity.com/paper/2303.09995/) (2023)
- [CSA-Trans: Code Structure Aware Transformer for AST](https://mlsanity.com/paper/2404.05767/) (2024)
- [DCMM-Transformer: Degree-Corrected Mixed-Membership Attention for Medical Imaging](https://mlsanity.com/paper/2511.12047/) (2025)

## Abstract (arXiv)

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 α-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 .