LG AI Research 291

Existing Transformer architectures show excellent performance, but have a limitation in that computational cost increases quadratically with sequence length. This makes them difficult to apply directly to long-sequence tasks such as document summarization or high-resolution image processing.

Existing Sparse Attention approaches have the drawback of applying the same modification to all layers, making it difficult to achieve the optimal balance between performance and cost. To address this, we propose SBM-Transformer, in which each attention head can adjust its sparsity on its own according to the data.

SBM-Transformer works by assigning a mixed-membership Stochastic Block Model (SBM) to each attention head.

• How it works: Each head treats input tokens as nodes and samples a Bipartite Graph connecting queries and keys.
• Graph sampling: Using an SBM parameterized by the input, it determines the query source nodes, key destination nodes, and the connection probabilities between clusters.
• Flexibility: Depending on the characteristics of the data, it flexibly switches between Dense attention and Sparse attention, optimizing computational cost.