Cluster-Former: Clustering-based Sparse Transformer for Question Answering

1 INTRODUCTION . Long-range contextual understanding has proven critical in many natural language processing ( NLP ) tasks . For example , the relevant context for correctly answering an open-domain question can arch over thousands of words . Encoding long sequences via deep neural networks , however , has remained an expensive and challenging task due to high demand on training time and GPU memory . Traditional sequence modeling methods ( Hochreiter & Schmidhuber , 1997 ) encode long sequences in a chronological order , which suffers high latency . In the place of sequential encoding , recent models such as Transformer ( Vaswani et al. , 2017 ) use simultaneous self-attention over the entire input instead , which has been successfully adopted in many NLP tasks such as textual entailment ( Devlin et al. , 2019 ) , dependency parsing ( Zhou & Zhao , 2019 ) , and summarization ( Lewis et al. , 2019 ) . A caveat with Transformer though is that building full connections over long sequences translates to quadratic growth on memory demand and computational complexity w.r.t . sequence length . One way to efficiently encode long sequences is to first chunk a sequence into much shorter ones with a sliding window , then build connections between the shorter sequences ( Figure 1 ( a ) ) . For example , Child et al . ( 2019 ) , Beltagy et al . ( 2020 ) and Zaheer et al . ( 2020 ) apply sparse attention to chunked sequences in hand-designed patterns in order to gather information from the chunks ( Figure 1 ( b ) ) . Choi et al . ( 2017 ) and Wang et al . ( 2019 ) first use a simpler model to filter chunked sequences , then process selected sequences with fully-connected self-attention . Rae et al . ( 2019 ) makes use of the shared memory of chunked sequences to build connections between them . However , these methods can not encode long-range dependencies with as much flexibility or accuracy as fully-connected selfattention , due to their dependency on hand-designed patterns . Recently , several studies ( Kitaev et al. , 2020 ; Tay et al. , 2020b ) propose to further improve the sparse attention mechanism by hashing or sorting the hidden states into different buckets ( Figure 1 ( c ) ) . These works mainly explore tasks with relatively short sequences , such as sentence-level Machine Translation ( MT ) , where the number of hashing vectors is relatively small ( less than 16 in Kitaev et al . ( 2020 ) ) , allowing randomly initialized hashing vectors to hash hidden states into correct buckets . However , how to use hashing-based attention in the context of long sequences ( e.g. , , up to thousands of words ) is still an unexplored territory . Our proposed framework for efficient long sequence encoding , Cluster-Former , marries both sliding-window and hashing-based methods to achieve effective local and long-range dependency encoding . Cluster-Former consists of two types of encoding layer . The first one ( noted as SlidingWindow Layer ) focuses on extracting local information within a sliding window . It applies Transformer to the hidden states of each chunked sequence independently , as shown in Figure 1 ( a ) . The other one ( noted as Cluster-Former Layer ) learns to encode global information beyond the initial chunked sequences . Specifically , we first apply clustering to the input hidden states so that similar hidden states are assigned to the same cluster , as shown in Figure 1 ( d ) . The clustered and sorted input is then divided uniformly into chunks , each encoded by a Transformer layer . Note that to make model training more efficient , the cluster centroids are not computed online but updated periodically ( every epoch or a few epochs ) . We accumulate the hidden states from the layer prior to the Cluster-Former layer in a memory bank , and apply the K-Means algorithm to form cluster centroids during each update cycle . Compared to previously discussed sparse attention based on pre-selected positions ( Figure 1 ( b ) ) or randomly-initialized hashing vectors ( Figure 1 ( c ) ) , experimental results show that our method can encode dependency across chunked sequences more effectively . Our contributions can be summarized as follows . ( i ) We propose Cluster-Former , a novel approach to capturing long-range dependencies more effectively than locality-sensitive hashing method . ( ii ) We propose a new Transformer-based framework to process long sequences by combining SlidingWindow and Cluster-Former layers to extract both local and global contextual information . ( iii ) Our model achieves the best performance on question answering datasets of Natural Questions ( long answer ) , SearchQA , and Quasar-T . 2 RELATED WORK . Efficient Transformers With Transformer models growing larger and larger , how to handle longer sequences arises as a critical challenge . Many works have been proposed to improve the computational and memory efficiency of Transformers , including Sparse Transformer ( Child et al. , 2019 ) , Routing Transformer ( Roy et al. , 2020 ) , Reformer ( Kitaev et al. , 2020 ) , Sinkhorn Transformer ( Tay et al. , 2020b ) , Longformer ( Beltagy et al. , 2020 ) , ETC ( Ainslie et al. , 2020 ) , Synthesizer ( Tay et al. , 2020a ) , Performer ( Choromanski et al. , 2020 ) , Linformer ( Wang et al. , 2020 ) , Linear Transformer ( Katharopoulos et al. , 2020 ) , and BigBird ( Zaheer et al. , 2020 ) . Tay et al . ( 2020c ) provided an excellent literature survey on this emerging topic . Our method falls into the setting of learnable sparse-attention patterns including Routing Transformer , Reformer and Sinkhorn Transformer . Our method is closer to Routing Transformer ( Roy et al. , 2020 ) which also uses cluster centroids to learn patterns , while we are targeting on quite different tasks ( language modeling VS question answering ) which leads to the significant difference on frameworks . Moreover , our cluster centroids are updated in very different ways ( online exponentially moving centroids VS periodical centroids update by KMeans ) . Long Sequence in Question Answering For tasks such as open-domain question answering ( Chen et al. , 2017 ) , a large volume of documents or paragraphs is usually retrieved to infer the answer , yielding extremely long context content . Despite the fact that state-of-the-art NLP models are capable of extracting answers amid complex context , they still struggle with extremely long input sequences . Recent advances that advocate the use of large-scale pre-trained models ( Lewis et al. , 2019 ; Liu et al. , 2019 ; Lan et al. , 2020 ) for question answering make this problem more prominent , due to tremendous memory consumption . To process long sequences , the most widely-used method is to first use a lightweight model to filter out redundant text , then use sliding-window-based approaches to encode the remaining sequences with a more sophisticated model . Chen et al . ( 2017 ) integrated bi-gram features into Information Retrieval ( IR ) methods to retrieve related documents more accurately . Wang et al . ( 2018 ) trained a paragraph selector using as the reward whether the entire system can obtain the correct answer or not . Lin et al . ( 2018 ) proposed to use a paragraph ranking model to curate data that are required for training reading comprehension models . Wang et al . ( 2019 ) trained a ranker to merge paragraphs for multi-passage reasoning . Asai et al . ( 2020 ) trained a recurrent retriever to select paragraphs for multi-hop question answering . Besides the above methods , directly applying Efficient Transformers to process long sequences in question answering is another option . In this paper , we focus on this direction by directly training our Cluster-Former on the long context without using lightweight model for context filtering . 3 PROPOSED APPROACH . The proposed framework to handle long sequences is pivoted on two types of Transformer layer : ( i ) Sliding-Window Layer ; and ( ii ) Cluster-Former Layer . The former focuses on encoding local sequence information , while the latter is on encoding global context and always built on top of the former layer . An overview of the two layers is illustrated in Figure 2 . 3.1 SLIDING-WINDOW LAYER . Despite that our focus is on capturing long-range dependencies for global context , local information also plays a critical role for knowledge propagation . Therefore , in the lower section of our network , we adopt the traditional sliding-window encoding mechanism . A sliding window segments a long sequence X into short , overlapping ones with window size l and stride m , as illustrated in Figure 2 ( a ) . Note that in this paper , we focus on question answering tasks , for which we concatenate the question Q with each sequence chunked from the document : H0k = [ Q ; X [ m× k : ( m× k + l ) ] ] , ( 1 ) where Q ∈ Rq×d denotes question embeddings given a QA task , q is the number of tokens in the question , and X ∈ Rx×d is the embeddings for all context , x is the number of tokens in context . k is the ID of the chunked sequence , l is the window size , and m is the stride of the sliding window . [ idx1 : idx2 ] indicates selecting rows between the index of idx1 and idx2 of the matrix . [ · ; · ] means concatenating the matrices along the row . We first use Transformer to encode each sequence in sliding window as follows : Hn+1k = Transformer ( H n k ) , ( 2 ) where Hn+1k ∈ R ( q+l ) ×d is the output of Transformer on the k-th sequence in the n-th layer . While it is not the final output of the n-th layer . As we expect the neighbouring sequences to share useful information in hidden states as well , we always set m < l to allow overlapping between sequences . We use the mean values of the Transformer hidden states at the overlapped tokens between windows as final outputs . To merge the representations from ( k − 1 ) -th sequence : Hn+1k [ q : q + l −m ] + = H n+1 k−1 [ q + m : end ] , Hn+1k [ q : q + l −m ] / = 2 , and merge representations from ( k + 1 ) -th sequence : Hn+1k [ q + m : end ] + = H n+1 k+1 [ q : q + l −m ] , Hn+1k [ q + m : end ] / = 2 , ( 3 ) where + = is to add matrices in-place and / = is to divide a matrix by a scalar value in-place . The merged hidden states Hn+1k ∈ R ( q+l ) ×d are the final outputs of the n-th layer . If the next layer is Cluster-Former , the output hidden states in this layer Hn+1k will be saved into memory bank for computing the cluster centroids . 3.2 CLUSTER-FORMER LAYER . Algorithm 1 Cluster Centroids Update 1 : Initialize Memory = Queue ( ) 2 : Centroids = GETCENTROIDS ( RandomVector ) 3 : 4 : function TRAIN ( Inputs ) 5 : for i = 1 , 2 , . . . , IterationNum do 6 : States = Sliding-Transformer ( Inputs [ i ] ) 7 : Memory.add ( States ) 8 : while len ( Memory ) > M do 9 : Memory.pop ( ) 10 : end while 11 : if i % ClusterUpdateFrequency == 0 then 12 : Centroids = GETCENTROIDS ( Memory ) 13 : end if 14 : Clusters = cluster States by Centroids 15 : States = Cluster-Former ( Clusters ) 16 : end for 17 : end function 18 : 19 : function GETCENTROIDS ( HiddenStates ) 20 : Centroids = K-Means ( HiddenStates ) 21 : Outputs = List ( ) 22 : Outputs [ 1 ] = Centroids [ 1 ] 23 : for i = 2 , 3 , . . . , ClusterNum do 24 : Outputs [ i ] = centroid from Centroids that is closest to Outputs [ i− 1 ] but not in Outputs 25 : end for 26 : return Outputs 27 : end function We introduce a Cluster-Former layer to add global representational power to Transformer beyond sliding windows . An in-depth visualization of the layer is illustrated in Figure 2 ( b ) . The input of the Cluster-Former layer comes from the hidden states of the prior layer ( in our case a Sliding-Window layer ) . After merging the overlaps between sequence chunks , the input of this layer is defined as : H̄n = [ Hn0 [ 0 : q + m ] ; ... ; H n k [ 0 : q + m ] ] , ( 4 ) where H̄n ∈ R ( qdx/me+x ) ×d is the hidden states to cluster , x is the number of tokens in the context . As the hidden states with larger cosine similarity are more likely to have higher attention weights , we build sparse selfattention only on the hidden states in the same cluster . In this work , we use KMeans as the chosen clustering method for simplicity . More advanced clustering algorithms have the potential of yielding better performance . Since running KMeans on the fly in each training iteration is computationally expensive , we decide to re-compute the cluster centroids with low frequency ( every epoch or a few epochs ) . In addition , to avoid dramatic changes in the cluster centroids due to limited hidden state inputs , we maintain a memory bank for the most recent hidden states . The entire procedure is depicted in Algorithm 1 . Once we compute the cluster centroids , we can directly use them for hidden state clustering as follows : vn = argmax ( Hn ( Cn ) T ||Hn||2||Cn||2 ) , ( 5 ) where Cn ∈ Rp×d are the cluster centroids for layer n , and p is the pre-defined number of clusters . The function argmax ( · ) performs on the last dimension and assigns all the input hidden states into different clusters based on the max value of cosine similarity between the hidden states and cluster centroids . vn ∈ R ( qdx/me+x ) is the assigned cluster IDs of all the input hidden states . Since the number of hidden states in different clusters can vary substantially , padding them to the maximum length for Transformer training will significantly increase the computational time . To make the extraction of global context more efficient , we greedily pick the cluster centroids based on the nearest neighbour ( measured by cosine similarity ) as shown in the function GETCENTROIDS in Algorithm 1 . Thus , the hidden states with similar cluster IDs are also close to each other . Then , we can directly sort the cluster IDs of hidden states and uniformly chunk the hidden states ( same window size and stride m ) : un = argsort ( vn ) , ank = u n [ mk : m ( k + 1 ) ] , Enk = H n [ ank ] , ( 6 ) where the function argsort ( · ) is to obtain the indexes of input values sorted in order ( same values sorted by the corresponding position of hidden states ) . ank ∈ Rm is the chunked indexes of the hidden states . Enk ∈ Rm×d is the k-th clustered hidden states , and we will run Transformer on top of it to build the connection beyond the words in the initial sliding window as follows : En+1k = Transformer ( E n k ) . ( 7 ) After updating the hidden states , we map them back to the order before clustering : H̄n+1 = [ En+10 ; E n+1 1 ; ... ; E n+1 K ] , ān = [ an0 ; a n 1 ; ... ; a n K ] , ( 8 ) H̄n+1 [ ān ] = clone ( H̄n+1 ) , ( 9 ) where H̄n+1 is the final output hidden state of this layer and has the same word order as the input H̄n . In experiments , we stack these two types of layer interchangeably to capture both global and local context efficiently . 4 EXPERIMENTS 4.1 DATASETS We evaluate our proposed approach on multiple question answering benchmarks . The statistics of all the datasets are summarized in Table 1 . Quasar-T1 ( Dhingra et al. , 2017 ) : The goal of this task is to answer open-domain questions from Trivia Challenge . All the passages harvested through information retrieval can be used to answer questions . The task requires the model to generate answers in phrases . The evaluation metric on this dataset is based on Exact Match and F1 score of the bag-of-words matching . Our evaluation tool2 comes from the SQuAD dataset . SearchQA3 ( Dunn et al. , 2017 ) : The setting of this dataset is the same as Quasar-T , except that the questions are sourced from Jeopardy ! instead . 1https : //github.com/bdhingra/quasar 2https : //rajpurkar.github.io/SQuAD-explorer/ 3https : //github.com/nyu-dl/dl4ir-searchQA Natural Questions4 ( Kwiatkowski et al. , 2019 ) : This task aims to answer questions based on a given Wikipedia document , and has two settings . ( i ) Long answer : select a paragraph that can answer the question based on the Wikipedia document if any . ( ii ) Short answer : extract an answer phrase from the document if the document contains the answer . As the given document may not contain answer , we can either predict an answer or predict no answer . The evaluation metric on this dataset is the F1 score , where true positives are exactly correct answers , false positives are incorrect answer predictions , and false negatives are incorrect “ no answer ” predictions . As the test set is hidden , we split 5 % of the training set for validation , and use the original validation set for testing . We use the official tool from the dataset to evaluate our models . We also submit our best model to the leaderboard .

Summary: Cluster-former is the latest proposal for enabling transformers to deal with long input sequences. Such sequences are particularly problematic for problems like question answering, QA, (or summarization), where the context can be arbitrarily long, and effectively open-ended when the setup includes a context retrieval component (e.g., as in OpenQA). Cluster-Former combines local information by encoding sequence chunks separately with a sliding window, then injects clustering layers, that use k-means to compute centroids to cluster hidden states and capture global information. The approach yields state-of-the-art, and top-of-leaderboard, results on Natural Questions (long answers).