Timezone: »

 
Poster
Transformers meet Stochastic Block Models: Attention with Data-Adaptive Sparsity and Cost
Sungjun Cho · Seonwoo Min · Jinwoo Kim · Moontae Lee · Honglak Lee · Seunghoon Hong

Tue Nov 29 02:00 PM -- 04:00 PM (PST) @ Hall J #419
in Poster Session 2 »
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.
Keywords: sparse attention · Stochastic Block Model · Efficient Transformers
[ Paper]  [ Poster]  [ OpenReview] 

Author Information

Sungjun Cho (LG AI Research)
Seonwoo Min (Seoul National University)
Jinwoo Kim (KAIST)
Moontae Lee (University of Illinois at Chicago)
Honglak Lee (LG AI Research / U. Michigan)
Seunghoon Hong (Korea Advanced Institute of Science and Technology)

More from the Same Authors