Long Document Ranking with Query-Directed Sparse Transformer (2020)

Abstract: The computing cost of transformer self-attention often necessitates breaking long documents to fit in pretrained models in document ranking tasks. In this paper, we design Query-Directed Sparse attention that induces IR-axiomatic structures in transformer self-attention. Our model, QDS-Transformer, enforces the principle properties desired in ranking: local contextualization, hierarchical representation, and query-oriented proximity matching, while it also enjoys efficiency from sparsity. Experiments on four fully supervised and few-shot TREC document ranking benchmarks demonstrate the consistent and robust advantage of QDS-Transformer over previous approaches, as they either retrofit long documents into BERT or use sparse attention without emphasizing IR principles. We further quantify the computing complexity and demonstrates that our sparse attention with TVM implementation is twice more efficient that the fully-connected self-attention. All source codes, trained model, and predictions of this work are available at https://github.com/hallogameboy/QDS-Transformer .

## introduction pre-trained transformers such as bert @xcite effectively transfer language understanding to better relevance estimation in many search ranking tasks @xcite @xcite . nevertheless, the effectiveness comes at the quadratic cost o(n 2 ) in computing complexity corresponds to the text length n, prohibiting its direct application to long documents. prior work adopts quick workarounds such as document truncation or splitting-and-pooling to retrofit the document ranking task to pretrained transformers. whilst there have been successes with careful architecture design, those bandit-solutions inevitably introduce information loss and create complicated system pipelines. intuitively, effective document ranking does not require fully connected self-attention between all query and document terms. the relevance matching between queries and documents often takes place at text segments as opposed to individual tokens @xcite , suggesting that a document term may not need information thousands of words away @xcite , and that not all document terms are useful to calculate the relevance to the query @xcite . the fully connected attention matrix includes many unlikely connections that create efficiency debt in computing, inference time, parameter size, and training convergence. this paper presents query-directed sparse transformer (qds-transformer) for long document ranking. in contrast to retrofitted solutions, qds-transformer fundamentally considers the desirable properties for assessing relevance by focusing on attention paths that matter. using sparse local attention @xcite , our model removes unnecessary connections between distant document tokens. using global attention upon sentence boundaries, our model further incorporates the hierarchical structures within documents. last but not the least, we use global attention on all query terms that direct the focus to the relevance matches between query-document term pairs. these three attention patterns in our query-directed sparse attention, as illustrated in figure 1 , permit global dissemination of ir-axiomatic information while keeping computation compact and essential. in our experiments with trec deep learning track @xcite and three more fewshot document ranking benchmarks @xcite , qds-transformer consistently improves the standard retrofitting bert ranking baselines (e.g., max-pooling on paragraphs) by 5% ndcg. it also shows gains over more recent transformer architectures that induces various sparse structures, including sparse transformer, longformer, and transformer-xh, as they were not designed to incorporate the essential information required in document ranking. in the meantime, we also thoroughly quantify the efficiency improvement from our query-directed sparsity, showing that with tvm support @xcite , different sparse attention patterns lead to variant training and inference speed up, and in general qds-transformer enjoys 200%+ speed up compared to vanilla bert on long documents. our visualization also shows interesting learned attention patterns in qds-transformer. similar to the observation on bert in nlp pipeline @xcite , in lower qds-transformer levels, the attention focuses more on learning the local interactions and document hierarchies, while in higher layers the model focuses more on relevance matching with the query terms. we also show examples that qds attention may center on the sole sentence that directly answers the query, or may span across several sentences that cover different aspects of the query, depending on the scope of the intent; this brings the advantage of better interpretability based on sparse attention. ## related work neural models have demonstrated significant advances across various ranking tasks @xcite . early approaches investigated diverse ways to capture relevance between queries and documents @xcite @xcite . and recently the state-of-the-art in many text ranking tasks has been taken by bert or other pretrained language models @xcite @xcite @xcite , when sufficient relevance labels are available for fine-tuning (e.g., on ms marco @xcite ). the improved effectiveness comes with the cost of computing efficiency with deep pretrained transformers, especially on long documents. this stimulates studies investigating ways to retrofit long documents to bert's maximum sequence length limits (512). a vanilla strategy is to truncate or split the documents: @xcite applied bert ranking on each passage segmented from the document independently and explored different ways to combine the passage ranking scores, using the score of the first passage (bert-firstp), the best passage (bert-maxp) (also studied in @xcite ), or the sum of all passage scores (bert-sump). more sophisticated approaches have also been developed to introduce structures to transformer attentions. transformer-xl employs recurrence on a sequence of text pieces @xcite , transformer-xh @xcite ) models a group of text sequences by linking them with extra hop attention paths, and transformer kernel long (tkl) @xcite uses a sliding window over the document terms and matches them with the query terms using matching kernels @xcite . on the efficiency front, @xcite proposed reformer that employed locality-sensitive hashing and reversible residual layers to improve the efficiency of transformers. @xcite introduced sparse transformers to reduce the quadratic complexity to o(l √ l) by applying sparse factorizations to the attention matrix, making the use of self-attention possible for extremely long sequences. subsequent work @xcite leverage a similar idea in a more adaptive way. combining local windowed attention with a task motivated global attention, @xcite presented longformer with an attention mechanism that scales linearly with sequence length. ## preliminaries on document ranking given a query q and a set of candidate documents @xmath0, the document ranking task is to produce the ranking score f (q, d) for each candidate global attention length-w attention window query attention local attention length-w attention window token embedding layer 1-st transformer layer l-th transformer layer (l 2)-th transformer layer (l 1)-th transformer layer length-w attention window [sos] [sos] [sos] ranking score f (q, d) document based on their relevance to the query. bert ranker. the standard way to leverage pretrained bert in document ranking is to concatenate the query and the document into one text sequence, feed it into bert layers, and then use a linear layer on top of the last layer's @xcite token @xcite : this bert ranker can be fine-tuned using relevance labels on (q, d) pairs, as simple as a classification task, and has achieved strong performances in various text ranking benchmarks @xcite . ## qds-transformer recent research has shown that with sufficient training and fully-connected self-attention, bert learns attention patterns that capture meaningful structures in language @xcite or for specific tasks @xcite . however, this is not yet the case in long document ranking as computing becomes the bottleneck. this section first presents how we overcome this bottleneck by injecting ir-specific inductive bias as sparse attention patterns. then we discuss the efficient implementation of sparse attention. ## experimental methodologies this section discusses our experimental settings. trec 2019 deep learning track benchmark. we evaluate qds-transformer based on the document ranking task from this recent trec benchmark (craswell et al., 2020), specifically using the reranking subtask to rerank top-100 bm25 retrieved documents. the official evaluation metric is ndcg@10 on the test set. we also report map on test and mrr@10 on the development set. ## evaluation results this section evaluates qds-transformer in its effectiveness, attention patterns, and efficiency. we also analyze the learned query-directed attention weights and show case studies. compared with sparse transformers and longformer-qa, qds-transformer provides more than 5% improvement in nearly all datasets. the best baseline is transformer-xh, which creates structural sparsity by breaking a document into segments and introduces effective extra-hop attentions to jointly model the relevance of those ## conclusions qds-transformer improves the efficiency and effectiveness of pretrained transformers in long document ranking using sparse attention structures. the sparsity is designed to capture the principal properties (ir-axioms) that are crucial for relevance modeling: local contextualization, document structures, and query-focused matching. in four trec document ranking tasks with variant settings, qds-transformer consistently outperforms competitive baselines that retrofit to bert or use sparse attention not designed for document ranking. our experiments demonstrate the promising future of joint optimization of structural domain knowledge and efficiency from sparsity, while its current form is somewhat at the infancy stage. our