[NeurIPS 2022] Transformers meet Stochastic Block Models: Attention with Data-Adaptive Sparsity and Cost

Introduction

The Transformer architecture has been the go-to method for encoding sequential data, due to its superior performance in various fields such as Natural Language Processing, Computer Vision, and Molecular Representation Learning. While its multi-head attention module enjoys high expressive capacity by computing the so-called attention score matrix (Figures 1 and 2) that captures how contextually important one token is to another, this comes at a computational cost that is quadratic to the sequence length. This prohibits direct application of Transformers to tasks that involve long sequences such as document summarization, high-resolution image processing, and protein sequence modeling.

Figure 1. Architecture of standard Transformer^[1,2]

Figure 2. Attention-computation in standard Transformer

While there exist methods that accelerate Transformers by leveraging sparse attention patterns or low-rank approximations, these approaches apply the same modification across all attention layers and heads, imposing an inductive bias too strong that often leads to sub-optimal cost vs. performance trade-offs in downstream tasks. Considering that many state-of-the-art systems utilize a mixture of dense and sparse attention (e.g. BERT and GPT-3), it is evident that developing an attention-module that can flexibly adjust itself between sparse and dense attention based on the given task would be beneficial towards better performance. Hence, we propose SBM-Transformer, a Transformer model that can data-adaptively choose its attention-sparsity as well as computational cost by endowing each attention head a mixed-membership Stochastic Block Model (SBM).

Our Method: SBM-Transformer

During the forward-step of SBM-Transformer, each attention head views the input tokens as a set of nodes and samples a bipartite graph M that connects querys to keys. Then, the attention score $q_i{k}_j^T$ is computed if and only if the corresponding edge $\left(i,j\right)$ exists in the sampled graph.

Figure 3. Forward-step of SBM-Transformer^[3]

For the graph sampling procedure, we use a Stochastic Block Model that is parameterized by the given input. Each SBM requires three non-negative parameters: the cluster memberships of the query-source nodes $\hat{\mathbf{Q}}$, the cluster memberships of the key-destination nodes $\hat{\mathbf{K}}$, and the inter-cluster connection probabilities $\hat{\mathbf{S}}$. For $\hat{\mathbf{S}}$, we represent clusters with learnable embeddings C and take its inner product with itself to obtain the strength of each inter-cluster connection. For node-memberships $\hat{\mathbf{Q}}$ and $\hat{\mathbf{K}}$, we process the query and key representations through an MLP, projecting them from the token-representation space onto the node-representation space, and then take the inner product with the cluster embeddings C. Figure 4 shows the full procedure of parameterizing the underlying SBM based on the input queries/keys and the cluster embeddings.

Figure 4. Parameterization of each SBM from query-key representations and cluster embeddings

Once we are given all three parameters for the SBM, we run the fastRG^[4] algorithm (Figure 5) to sample the graph. The overall steps are: 1) normalize the parameters such that the node-memberships become probability distributions across a set of nodes (Lines 1-3). 2) Sample the number of edges to generate from a Poisson distribution (Line 4). 3) For each edge, sample a pair of clusters, and sample the source and destination node from the corresponding node-probability distributions (Lines 6-11). Note that the entire process takes cost linear to the number of edges, and the for-loop can run in parallel for further acceleration as each edge is sampled independently.

Figure 5. The fastRG Algorithm used to sample from the SBM^[3]

Despite its simplicity and efficiency, the graph sampling step is naturally discrete and hence naïve backpropagation cannot learn proper parameterizations for the SBM. To cope with this non-differentiability, we use a Straight-Through Estimator (STE) to pass the gradient past the graph sampling step and directly to the edge probabilities. This way, the probabilities of sampled edges can be adjusted in an end-to-end fashion as the gradients provide information on which edges were useful towards correct predictions. As a result, the backward step of SBM-Transformer also takes cost that is linear to the number of sampled edges, which is chosen by the model data-adaptively. The overall pipeline is depicted in Figure 6.

Figure 6. Overall attention module in SBM-Transformer^[3]

Figure 7. Examples attention patterns that can be represented by Stochastic Block Models^[3]

Surprisingly, we find that SBMs can express a wide variety of attention masks as shown in Figure 7, based on the placement of node- and cluster-representations in the latent feature space. The overall density of the attention mask can easily range from full attention to no attention, depending on whether the embeddings are gathered up closely or not. This flexibility of the SBM allows us to prove that SBM-Transformer preserves the same universal approximability of the original Transformer, despite assuming a low-rank structure on the latent graph space.

Experiments

For quantitative evaluation, we first test our proposed SBM-Transformer against vanilla Transformer as well as existing efficient Transformer variants on the Long Range Arena (LRA) benchmark^[5]. Results in Table 1 show that our SBM-Transformer model outperforms all efficient Transformer variants, and even the original Transformer with full-attention while using much less attention. Table 2 also shows that SBM-T achieves significant reduction in number of FLOPs as well as peak memory use during inference.

Table 1. Accuracy results on the LRA benchmark. For our SBM-T results, λ denotes the attention density regularizer that penalizes each sampled edge and the percentages in parentheses denote the attention density during test time^[3]

Table 2. Per-example relative FLOP counts and peak memory use during LRA test time^[3]

To test our model on the downstream NLP setting, we also perform evaluation on the General Language Understanding Evaluation (GLUE) benchmark^[6]. Table 3 shows the classification accuracies for each task, from which we can see that our model performs competitively against existing baselines.

Table 3. Accuracy results on the GLUE benchmark^[3]

For qualitative analysis, we visualize the per-example attention heatmap from Image and Pathfinder tasks in LRA to see which examples use dense or sparse attention for correct predictions. For LRA Pathfinder, where the task is to determine whether two dots are connected with a dashed line, we interestingly find that the way SBM-Transformer distributes attention aligns well with human perception: it uses dense attention to examples that are difficult for the human-eye, while using much less attention for the relatively easier ones. Performing the same analysis in LRA Image task, which is equivalent to CIFAR-10, we find that the model mainly uses attention to recognize the large shift in contrast between the object in the image and the background towards correct classification.

Figure 8. Examples Attention Density Maps from LRA Image (Left) and LRA Pathfinder (Right)^[3]

Conclusion

SBM-Transformer is an efficient variant of the widely-used Transformer architecture that can data-adaptively choose its attention sparsity between sparse and full attention without explicitly computing the full attention score matrix, thereby avoiding the quadratic cost when possible. We showed that the use of low-rank latent graph structures via SBMs allow SBM-Transformer to retain the expressive power of full-attention, with experiments on the LRA and GLUE benchmarks demonstrating its competitive performance against the original Transformer as well as other baselines. For future work, we hope to incorporate more GPU-friendly tensor operations for unstructured and fine-grained sparsity towards better optimization and further investigate the latent geometry of query-key interactions induced by SBM-Transformer during inference.

▶Transformers meet Stochastic Blockmodels: Attention with Data-Adaptive Sparsity and Cost (Link)