C luster F ormer: Neural Clustering Attention for Efficient and Effective Transformer (2022)

Abstract: Recently, a lot of research has been carried out to improve the efficiency of Transformer. Among them, the sparse pattern-based method is an important branch of efficient Transformers. However, some existing sparse methods usually use fixed patterns to select words, without considering similarities between words. Other sparse methods use clustering patterns to select words, but the clustering process is separate from the training process of the target task, which causes a decrease in effectiveness. To address these limitations, we design a neural clustering method, which can be seamlessly integrated into the Self-Attention Mechanism in Transformer. The clustering task and the target task are jointly trained and optimized to benefit each other, leading to significant effectiveness improvement. In addition, our method groups the words with strong dependencies into the same cluster and performs the attention mechanism for each cluster independently, which improves the efficiency. We verified our method on machine translation, text classification, natural language inference, and text matching tasks. Experimental results show that our method outperforms two typical sparse attention methods, Reformer and Routing Transformer while having a comparable or even better time and memory efficiency.

## introduction transformer @xcite has been widely used and achieved state-of-the-art results in a variety of nlp tasks such as neural machine translation @xcite , text classification, etc. its good effectiveness benefits from its core component self-attention mechanism which can capture global dependencies well. however, the large calculation and memory cost limit the further application of transformer on long sequence * the first two authors contributed equally † corresponding author tasks due to the complexity of o(n 2 d) of self-attention. as a result, many research works have been carried out to improve the efficiency of transformer @xcite . these efficient transformers can be roughly divided into two categories: approximation-based @xcite and sparse patternbased methods @xcite @xcite . regarding approximation-based methods, some works are based on low-rank approximation @xcite while others are based on kernels @xcite . specifically, linformer @xcite adopts a low-rank approximation idea and projects the length dimension of keys and values to a lower dimension (@xmath0. this reduces the complexity to o(n kd). however, the projection matrix in this method requires that all input sequences must be filled to the same length n , which makes it cannot handle the variable-length sequence well. linear transformer @xcite uses kernels and the associative property of matrix multiplication to linearize the softmax attention, which reduces the complexity to o(n d 2 ). however, the approximation error to the softmax matrix in self-attention can be large in some cases @xcite . sparse pattern-based methods introduce sparse patterns into the attention mechanism and limit the number of key vectors that the query vector should pay attention to. prior work @xcite @xcite @xcite proposed to use fixed sparse patterns to improve the efficiency of self-attention. for example, sparse transformer @xcite restricts the query to focus only on keys that are nearby or at fixed intervals. such fixed sparse patterns do not consider the similarity between the query and different keys, and directly filter keys according to their location, which re-sults in a degradation of the model effectiveness. more recently, the clustering sparse patterns are proposed. reformer @xcite and routing transformer @xcite use locality sensitive hashing (lsh) and k-means algorithms, respectively, to divide the words in the sequence into different clusters, and then the attention operation is restricted within each cluster independently. in this way, they reduce the complexity to o(n log n d) and o(n √ n d), respectively. however, in reformer and routing transformer, both lsh and k-means only play the role of cluster partitioning, but run separately from the attention network training. in addition, these two methods also have the problem of inconsistency in the similarity measure between the clustering and attention operation. lsh and k-means respectively use the hash value obtained by random projections and the negative euclidean distance as the similarity measure between the input vectors while the attention operations use the inner product. therefore, such a sparse pattern idea often results in the reduced effectiveness. to address the reduced effectiveness issue of many efficient transformers, especially sparse pattern-based methods, we propose neural clustering method to learn the sparse pattern of the attention. it can be seamlessly integrated into neural networks for joint training and optimization. in our method, the cluster center (centroid) is updated by a weighted sum of all word hidden states. at the same time, the members of clusters are divided according to the subordinate matrix of the centroids and word hidden states. the optimization of the clustering loss can guide the representation of word hidden states while learning cluster centroids. the integration of the neural clustering method and attention training enables our neural clustering attention to perform better than previous clustering-based sparse attention mechanisms. our neural clustering method is a general clustering method, in the sense that in addition to being integrated into the network of the specific task, it is can handle clustering tasks alone. our overall model is called clusterformer and it is obtained by replacing the self-attention mechanism in transformer with neural clustering attention mechanism. in order to validate the benefits of clusterformer, we have carried out comparison experiments of the efficiency and effectiveness respectively. for efficiency, we provide a detailed analysis about the time and memory on dataset 20news of text classification. results show that our model has a comparable or even better efficiency compared with two typical sparse attention models, reformer and routing transformer. especially, when the sequence length exceeds 2000, our model, routing transfomer and reformer reduce the memory of transformer by 53.8%, 60.8% and 31.8% while reducing the training time by 51.4%, 41.8% and 14.4%, respectively on gpu. for effectiveness, we test it on machine translation, text classification, natural language inference, and text matching tasks. experimental results show that on all tasks, our model consistently outperforms reformer and routing transformer. in particular, our method improves the accuracy by 15.6% and 7.2% on scitail datasets of natural language inference task compared with reformer and routing transformer, respectively. the major contributions of our work are as follows: ## model in this section, we first introduce our neural clustering method. then, we introduce our neural clustering attention mechanism which combines our clustering method and the self-attention mechanism. ## experiments in order to verify the effectiveness and efficiency of our method, we carried out the following tasks. we choose transformer and its clustering-pattern variants (reformer, routing transformer) as baseline models. the implementations of the attention layer of reformer and routing transformer refer to the open source codes 12 . for a fair comparison, our proposed method and baseline models have the same architecture, except for the attention layer. ## conclusion in this paper, we propose a neural clustering attention mechanism to address the reduced effectiveness issue in sparse attention methods. this issue is mainly caused by the introduction of a sparse pattern that is separated from the target task or does not consider the similarity between words. in our method, we design a neural clustering algorithm to better capture critical pairs of dependencies. we integrate this clustering algorithm and the neural network to jointly train and optimize with specific tasks together to further contribute to the effectiveness and efficiency. the experimental results show that our model can achieve better effectiveness and a comparable or even better efficiency, compared with the latest typical sparse attention models, reformer and routing transformer.