Optimizing Guided Traversal for Fast Learned Sparse Retrieval

Document retrieval for searching a large dataset often uses a sparse representation of document feature vectors implemented as an inverted index which associating each search term with a list of documents containing such a term. Recently learned sparse representations have been developed to compute term weights using a neural model such as transformer based retriever (Dai and Callan, 2020; Bai et al., 2020; Formal et al., 2021b, a; Mallia et al., 2021; Lin and Ma, 2021) and deliver strong relevance results, together with document expansion (e.g.  (Cheriton, 2019)). A downside is that top $k$ document retrieval latency using a learned sparse representation is much large than using the BM25 model as discussed in  (Mallia et al., 2021; Mackenzie et al., 2021). In the traditional BM25-based document retrieval with additive ranking, a dynamic index pruning strategy based on top $k$ threshold is very effective by computing the rank score upper bound on the fly for each visited document during index traversal in order to skip low-scoring documents that are unable to appear in the final top $k$ list. Well known traversal algorithms with such dynamic pruning strategies include MaxScore (Turtle and Flood, 1995) and WAND (Broder et al., 2003), and their block-based versions Block-Max WAND (BMW) (Ding and Suel, 2011) and Block-Max MaxScore (BMM) (Chakrabarti et al., 2011; Dimopoulos et al., 2013).

Mallia et al. (Mallia et al., 2022) propose to use BM25 to guide traversal, called GT, for fast learned sparse retrieval because the distribution of learned weights results in less pruning opportunities and they conducted an evaluation with retrieval model DeepImpact (Mallia et al., 2021). One variation they propose is to compute the final rank scoring as a linear combination of the learned weights and BM25 weights, denoted as GTI. GT is a special case of GTI and this paper treats GTI as the main baseline. Since the BM25 weight for a document term pair may not exist in a learned sparse index, zero filling is used in Mallia et al. (Mallia et al., 2022) to align the BM25 and learned weight models. During our evaluation using GT for SPLADE v2 and its revision SPLADE++ (Formal et al., 2021a, 2022), we find that as retrieval depth $k$ decreases, BM25 driven skipping becomes too aggressive in dropping documents desired by top $k$ ranking based on learned term weights, which can cause a significant relevance degradation. In addition, there is still some room to further improve index alignment of GTI for more accurate BM25 driven pruning.

To address the above issues, we improve our earlier pruning study on dual guidance with combined BM25 and learned weights (Qiao et al., 2022). Our work generalizes GTI by constraining the pruning influence of BM25 and providing an alternative smoothing method to align the BM25 index with learned weights. In Section 4, we propose a two-level parameterized guidance scheme with index alignment, called 2GTI, to manage pruning decisions during MaxScore based traversal. We analyze some formal properties of 2GTI on its relevance behaviors and configuration conditions when 2GTI outperforms a two-stage top $k$ search algorithm for a query in relevance.

Section 5 and Appendix A present an evaluation of 2GTI with SPLADE++ (Formal et al., 2021b, a, 2022) and uniCOIL (Lin and Ma, 2021; Gao et al., 2021a) in addition to DeepImpact (Mallia et al., 2021) when using MaxScore on the MS MARCO datasets. This evaluation shows that when retrieval depth $k$ is small, or when the BM25 index is not well aligned with the underlying learned sparse representation, 2GTI can outperform GTI and retain relevance more effectively. In some cases, there is a tradeoff that 2GTI based retrieval may be slower than that of GTI while 2GTI is still much faster than the original MaxScore method without BM25 guidance. 2GTI is also effective for the BEIR datasets in terms of the zero-shot relevance and retrieval latency. In Appendix B, we have extended the use of 2GTI for a BMW-based algorithm such as VBMW (Mallia et al., 2017). We demonstrate that 2GTI with VBMW can be useful for a class of short queries and when $k$ is small.

The top-$k$ document retrieval problem identifies top ranked results in matching a query. A document representation uses a feature vector to capture the semantics of a document. If these vectors contain much more zeros than non-zero entries, then such a representation is considered sparse. For a large dataset, document retrieval often uses a simple additive formula as the first stage of search and it computes the rank score of each document $d$ as:

where $Q$ is the set of all search terms, $w(t,d)$ is a weight contribution of term $t$ in document $d$, and $w_{t}$ is a document-independent or query-specific term weight. Assume that $w(t,d)$ can be statically or dynamically scaled, this paper views $w_{t}=1$ for simplicity of presentation. An example of such formula is BM25 (Jones et al., 2000) which is widely used. For a sparse representation, a retrieval algorithm often uses an inverted index with a set of terms, and a document posting list of each term. A posting record in this list contains document ID and its weight for the corresponding term.

Threshold-based skipping. During the traversal of posting lists in document retrieval, the previous studies have advocated dynamic pruning strategies to skip low-scoring documents, which cannot appear on the final top-$k$ list  (Broder et al., 2003; Strohman and Croft, 2007). To skip the scoring of a document, a pruning strategy computes the upper bound rank score of a candidate document $d$, referred to as $Bound(d)$.

If $Bound(d)\leq\theta$ where $\theta$ is the rank score threshold in the top final $k$ list, this document can be skipped. For example, WAND (Broder et al., 2003) uses the maximum term weights of documents of each posting list to determine the rank score upper bound of a pivot document while BMW (Ding and Suel, 2011) and its variants (e.g.  (Mallia et al., 2017)) optimize WAND use block-based maximum weights to compute the score upper bounds. MaxScore (Turtle and Flood, 1995) uses term partitioning and the top-$k$ threshold to skip unnecessary index visitation and scoring computation. Previous work has also pursued a “rank-unsafe” skipping strategy by deliberately over-estimating the current top-$k$ threshold by a factor (Broder et al., 2003; Macdonald et al., 2012; Tonellotto et al., 2013; Crane et al., 2017).

Learned sparse representations. Earlier sparse representation studies are conducted in (Zamani et al., 2018), DeepCT (Dai and Callan, 2020), and SparTerm (Bai et al., 2020). Recent work on this subject includes SPLADE (Formal et al., 2021b, a, 2022), which learns token importance for document expansion with sparsity control. DeepImpact (Mallia et al., 2021) learns neural term weights on documents expanded by DocT5Query (Cheriton, 2019). Similarly, uniCOIL (Lin and Ma, 2021) extends the work of COIL (Gao et al., 2021a) for contextualized term weights. Document retrieval with term weights learned from a transformer has been found slow in  (Mallia et al., 2022; Mackenzie et al., 2021). Mallia et al. (Mallia et al., 2022) state that the MaxScore retrieval algorithm does not efficiently exploit the DeepImpact scores. Mackenzie et al. (Mackenzie et al., 2021) view that the learned sparse term weights are “wacky” as they affect document skipping during retrieval thus they advocate ranking approximation with score-at-a-time traversal.

Our scheme uses a hybrid combination of BM25 and learned term weights, motivated by the previous work in composing lexical and neural ranking (Yang et al., 2022; Lin et al., 2021b; Gao et al., 2021b; Ma et al., 2021; Li et al., 2022). GTI adopts that for final ranking. A key difference in our work is that hybrid scoring is used for two-level pruning control and its formula can be different from final ranking. The multi-level hybrid scoring difference provides an opportunity for additional pruning and its quality control. Thus the outcome of 2GTI is not a simple linear ranking combination of BM25 and learned weights and two-level guided pruning yields a non-linear ensemble effect to improve time efficiency while retaining relevance. Our evaluation will include a relevance and efficiency comparison with MaxScore using a simple linear combination.

This paper mainly focuses on MaxScore because it has been shown more effective for relatively longer queries (Mallia et al., 2019a). We also consider VBMW (Mallia et al., 2017) because it is generally acknowledged to represent the state of the art (Mackenzie et al., 2021) for many cases, especially when $k$ is small and the query length is short (Mallia et al., 2019a).

Figure  1 shows the performance of the original MaxScore retrieval algorithm without BM25 guidance, GTI, and the proposed 2GTI scheme in terms of MRR@10 and recall@k when varying top $k$ in searching MS MARCO passages on Dev query set. Here $k$ is the targeted number of top documents to retrieve and it is also called retrieval depth sometime in the literature. Section 5 has more detailed dataset and index information. For both SPLADE++ and uniCOIL, we build the BM25 model following  (Mallia et al., 2022) to expand passages first using DocT5Query, and then use the BERT’s Word Piece tokenizer to tokenize the text, and align the token choices of BM25 with these learned models. From Figure 1, there are significant recall and MRR drops with GTI when $k$ varies from 1,000 to 10. There are two reasons contributing to the relevance drops.

1. (1)

   When the number of top documents $k$ is relatively small, the relevance drops significantly. As $k$ is small, dynamically-updated top $k$ score threshold becomes closer to the maximum rank score of the best document. Fewer documents fall into top $k$ positions and more documents below the updated top $k$ score threshold would be removed earlier. Then the accuracy of skipping becomes more sensitive. The discrepancy of BM25 scoring and learned weight scoring can cause good candidates to be removed inappropriately by BM25 guided pruning, which can lead to a significant relevance drop for small $k$.

2. (2)

   The relevance drop for SPLADE++ with BM25 guided pruning is noticeably much more significant than uniCOIL. That can be related to the fact that SPLADE++ expands tokens of each document tokens differently and much more aggressively than uniCOIL. As a result, 98.6% of term document pairs in SPLADE++ index does not exist in the BM25 index even after docT5Query document expansion while this number is 1.4% for uniCOIL. Thus, BM25 guidance can become less accurate and improperly skip more good documents.

With the above consideration, our objective is to control the influence of BM25 weights in a constrained manner for safeguarding relevance prudently, and to develop better weight alignment when the BM25 index is not well aligned with the learned sparse index. In Figure 1, the recall@k number of 2GTI marked with blue squares is similar to that of the original method without BM25 guidance. Their MRR@10 numbers overlapped with each other, forming a nearly-flat lines, which indicates their MRR@10 numbers are similar even $k$ decreases. The following two sections present our solutions in addressing the above two issues respectively.

In producing an inverted index, all words use lower case letters. Following GT, we packed the learned score and the term frequency in the same integer. For DeepImpact, we adopt GT’s index¹¹1https://github.com/DI4IR/dual-score directly. The BM25-T5’s index is dumped from the DeepImpact index. Both BM25-T5 and DeepImpact are using natural words tokenization.

SPLADE and uniCOIL use the BERT’s Word Piece tokenizer. In order to align with them, the BM25-T5-B index reported in the following tables uses the same tokenizer as well. The impact scores of uniCOIL is obtained from Pyserini (Lin et al., 2021a) ²²2https://github.com/castorini/pyserini/blob/master/docs/experiments-unicoil.md. For SPLADE, in order to achieve the best performance, we retrained the model following the setup in SPLADE++ (Formal et al., 2022). We start from the pretrained model coCondenser (coCondenser, 2021) and distill using the sentenceBERT hard negatives ³³3https://huggingface.co/datasets/sentence-transformers/msmarco-hard-negatives from a cross-encoder teacher (MiniLM-L-6-v2, 2022) with MarginMSE loss. For FLOP regularization, we use 0.01 and 0.008 for query and documents respectively. We construct the inverted indexes, convert them to the PISA format, and compress them using SIMD-BP128 (Lemire and Boytsov, 2015) following (Mallia et al., 2019a, 2022).

Table 1 shows the dataset and index characteristics of the different weighting models on the MS MARCO Dev dataset. Following (Mackenzie et al., 2021), we assume that a query can be pre-processed with a ”pseudo-document” trick that assigns custom weights to query terms in uniCOIL and SPLADE. Therefore, there may be token repetition in each query to reflect token weighting. Column 1 is the mean query length in tokens without or with counting duplicates. Column 3 is the inverted index size while the last column is the size after merging BM25 and learned weights in the index.

The C++ implementation of 2GTI with the modified MaxScore and VBMW algorithms are embedded in PISA (Mallia et al., 2019b), and the code will be released in https://github.com/Qiaoyf96/2GTI. Our evaluation using this implementation runs as a single thread on a Linux server with Intel i5-8259U 2.3GHz and 32GB memory. Weights are chosen by sampling queries from the MS MARCO training dataset.

Metrics. For MS MARCO Dev set, we report the relevance in terms of mean reciprocal rank (MRR@10 on passages and MRR@100 on documents), following the official leader-board standard. We also report the recall@k ratio which is the percentage of relevant-labeled results appeared in the final top-$k$ results. For TREC DL test sets, we report normalized discounted cumulative gain (nDCG@10) (Järvelin and Kekäläinen, 2002). The above reporting follows the common practice of the previous work (e.g.  (Mallia et al., 2021; Gao et al., 2021a, b; Formal et al., 2021a)).

Before timing queries, all compressed posting lists and metadata for tested queries are pre-loaded into memory, following the same assumption in (Khattab et al., 2020; Mallia et al., 2017). Retrieval mean response times (MRT) are reported in milliseconds. The 99th percentile time ($P_{99}$) is reported within parentheses in the tables below, corresponding to the time occurring in the 99th percentile denoted as tail latency in (Mackenzie et al., 2018).

Statistical significance. For the reported numbers on MS MARCO passage and document Dev sets in the rest of this section, we have performed a pairwise t-test on relevance difference between 2GTI and a GTI baseline, and between 2GTI and the original learned sparse retrieval without BM25 guidance. No statistically significant degradation has been observed at the 95% confidence level. We have also performed a pairwise t-test comparing the reported relevance numbers of 2GTI and GTI and mark ‘^†’ in the evaluation tables if there is a statistically significant improvement by 2GTI over GTI at the 95% confidence level. We do not perform a t-test on DL’19 and DL’20 query sets as the number of queries in these sets is small.

Overall results with MS MACRO. Table 2 lists a comparison of 2GTI with the baseline using three sparse representations for retrieval on MS MARCO and TREC DL datasets. 2GTI uses scaled filling alignment as default while GTI uses zero filling as specified in (Mallia et al., 2022). The $\gamma$ value is chosen the same for GTI and 2GTI for each representation, which is the best for most of cases. The “accurate” configuration denotes the one that reaches the highest relevance score. The “fast” configuration denotes the one that reaches a relevance score within 1% of the accurate configuration while being much more faster.

2GTI vs. GTI in SPLADE++. Table 2 shows 2GTI with default scaled filling significantly outperforms GTI with default zero filling for SPLADE++, where BM25 index is not well aligned. “SPLADE++-Org” denotes the original MaxScore retrieval performance using SPLADE++ model trained by ourselves and its MRR@10 number is higher than what has been reported in  (Formal et al., 2022). When $k=1,000$, GT is slightly better than GTI, and with the fast configuration, MRR@10 of 2GTI is 32.4% higher than that of GT while 2GTI is 7.8x faster than GT for the Dev set. The significant increase in nDCG@10 and decrease in the MRT are also observed in DL’19 and DL’20. When $k=10$, there is also a large relevance increase and time reduction from GTI or GT to 2GTI for all three test sets. For example, the relevance is 46.4% or 44.6% higher and the mean latency is 5.2x or 5.3x faster for the Dev set.

Compared to the original MaxScore method, 2GTI has about the same relevance score for both $k=10$ and $k=1,000$ while having much smaller latency. For example, 6.5x reduction (278ms vs. 43.1ms) for the Dev passage set when $k=1,000$ and 5.3x reduction when $k=10$ (121ms vs 22.7ms) with the 2GTI-fast configuration.

2GTI vs. GTI in DeepImpact and uniCOIL. As shown in Table 2, GTI (or GT) performs very well for $k=1,000$ in both DeepImpact and uniCOIL in speeding up retrieval while maintaining a relevance similar as the original retrieval. The two-level differentiation for dynamic index pruning does not improve relevance or shorten retrieval time. This can be explained as BM25-T5 index is well aligned with the DeepImpact index and with the uniCOIL index. Also because of this reason, filling to address index alignment is not needed with no improvement in these two cases.

When $k$ decreases from 1,000 to 10, as shown in Figure 1 discussed in Section 3, the recall ratio starts to drop, and relevance effectiveness degrades. When $k=10$ as shown in Table 2, DeepImpact-2GTI-fast can increase MRR@10 from 0.3375 by GTI to 0.3395 for the Dev set and deliver slightly higher MRR@10 or nDCG@10 scores than GTI in DL’19 and DL’19 sets. For uniCOIL, 2GTI-fast increases MRR@10 from 0.3384 by GTI to 0.3548 for the Dev set and increases nDCG@10 from 0.6959 to 0.7135 for DL’19. There is also a modest relevance increase for DL’20 passages with $k=10$ and a similar trend is observed for the document retrieval task. The price paid for 2GTI is its retrieval latency increase while its latency is still much smaller than the original retrieval time.

Design options with weight alignment and threshold over-estimation. Table 3 examines the impact of weight alignment and a design alternative based on threshold over-estimation for MS MARCO passage Dev set using SPLADE++ when $k=10$. In the top portion of this table, threshold over-estimation by a factor of $F$ (1.1, 1.3, and 1.5) is used in the original retrieval algorithm without BM25 guidance, and these factor choices are similar as ones in  (Macdonald et al., 2012; Tonellotto et al., 2013; Crane et al., 2017). That essentially sets $\alpha=0$, $\beta=0$, and $\gamma=0$ while multiplying $\theta_{Gl}$ and $\theta_{Lo}$ by the above factor in 2GTI. The result shows that even threshold over-estimation can reduce the retrieval time, relevance reduction is significant, meaning that the aggressive threshold used causes incorrect dropping of some desired documents.

The second portion of Table 3 examines the impact of different weight filling methods described in Section 4.3 for alignment when they are applied to GTI and 2GTI, respectively. In both cases, scaled filling marked as “/s” is most effective while one-filling marked as “/1” outperforms zero-filling marked as “/0” also. The MRT of 2GTI/s becomes 10.5x smaller than 2GTI/0 while there is no negative impact to its MRR@10. The MRT of GTI/s is about 13.0x smaller than GTI/0 while there is a large MRR@10 number increase.

A validation on 2GTI’s properties. To corroborate the competitiveness analysis in Section 4.2, Table 4 gives MRR@10 scores and retrieval times in milliseconds of the algorithms with different configurations on the Dev set of MS MARCO passages with $k=10$ and SPLADE++ weights. The result shows that the listed configurations of 2GTI have a higher MRR@10 number than 2-stage search $R2_{\alpha,\gamma}$, and 2GTI with $\alpha=\beta=1$ that behaves as GTI. MRR@10 of ranking with a simple linear combination of BM25 and learned weights is only slightly higher than 2GTI, but it is much slower.

Sensitivity on weight distribution. We have distorted the SPLADE++ weight distribution in several ways to examine the sensitivity of 2GTI and found that 2GTI is still effective. For example, we apply a square root function to the neural weight of every token in MS MARCO passages, the relevance score of both original retrieval and 2GTI drops to 0.356 MRR@10 due to weight distortion, while 2GTI is 5.0x faster than the original MaxScore when $k=10$.

Efficient SPLADE model. Table 5 shows the application of 2GTI in a recently published efficient SPLADE model (Lassance and Clinchant, 2022) which has made several improvements in retrieval speed. We have used the released checkpoint of this efficient model called BT-SPLADE-L, which has a weaker MRR@10 score, but significantly faster than our trained SPLADE baseline reported in Table 2. When used with this new SPLADE model, 2GTI/s-Fast version results in a 2.2x retrieval time speedup over MaxScore. Its MRR@10 is higher than GTI/s and has less than 1% degradation compared to the original MaxScore.

The contribution of this paper is a two-level parameterized guidance scheme with index alignment to optimize retrieval traversal with a learned sparse representation. Our formal analysis shows that a properly configured 2GTI algorithm including GTI can outperform a two-stage retrieval and re-ranking algorithm in relevance.

Our evaluation shows that the proposed 2GTI scheme can make the BM25 pruning guidance more accurate to retain the relevance. For MaxScore with SPLADE++ on MS MARCO passages, 2GTI can lift relevance by up-to 32.4% and is 7.8x faster than GTI when $k=1,000$, and by up-to 46.4% more accurate and 5.2x faster when $k=10$. In all evaluated cases, 2GTI is much faster than the original retrieval without BM25 guidance. For example, up-to 6.5x faster than MaxScore on SPLADE++ when $k=10$. We have also observed similar performance patterns on BEIR datasets when comparing 2GTI with GTI and the original MaxScore using SPLADE++ learned weights. Compared to other options such as threshold underestimation to reduce the influence of BM25 weights, the two-level control is more accurate in maintaining the strong relevance with a much lower time cost. While our study is mainly centered with MaxScore-based retrieval, 2GTI can be used for VBMW and our evaluation shows that VBMW-2GTI can be a preferred choice for a class of short queries without stop words when $k$ is small.

Acknowledgments. We thank Wentai Xie, Xiyue Wang, Tianbo Xiong, and anonymous referees for their valuable comments and/or help. This work is supported in part by NSF IIS-2225942 and has used the computing resource of the XSEDE and ACCESS programs supported by NSF. Any opinions, findings, conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the NSF.