LG AI Research 293

The Multi-Head Attention of the existing Transformer incurs computational cost proportional to the square of the sequence length, which limits its ability to handle long input sequences. Existing Sparse attention approaches have suffered from a trade-off between performance and cost due to excessive Inductive bias.

To address this, the proposed SBM-Transformer assigns a mixed-membership Stochastic Block Model (SBM) to each attention head. This allows it to flexibly adjust between Dense and Sparse attention according to the characteristics of the data, optimizing computational cost.

The core operating mechanism of SBM-Transformer is as follows:

• Graph sampling: Each attention head treats input tokens as nodes and samples a bipartite graph connecting Query and Key.
• Parameterization: Query/Key representations are mapped into a node representation space via an MLP, then dot-producted with Cluster embeddings to determine cluster membership.
• fastRG algorithm: The fastRG algorithm is used at the sampling stage to efficiently generate graphs, and this process can be parallelized for acceleration.

In addition, to address the Non-differentiability problem of graph sampling, the Straight-Through Estimator (STE) is utilized to enable Backpropagation.