BM25S: Orders of magnitude faster lexical search via eager sparse scoring

Many variants of BM25 exist, which could lead to significant confusion about the exact scoring method used in a given implementation Kamphuis et al. (2020). By default, we use the scoring method proposed by Lucene. Thus, for a given query $Q$ (tokenized into $q_{1},\ldots,q_{|Q|}$) and document $D$ from collection $C$, we compute the following score⁸⁸8We follow notations by Kamphuis et al. (2020):

	$\displaystyle B(Q,D)$	$\displaystyle=\sum_{i=1}^{|Q|}S(q_{i},D)$
		$\displaystyle=\sum_{i=1}^{|Q|}\text{IDF}(q_{i},C)\frac{\text{TF}(q_{i},D)}{\mathcal{D}}$

where $\mathcal{D}=\text{TF}(t,D)+k_{1}\left(1-b+b\frac{|D|}{L_{avg}}\right)$, $L_{avg}$ is the average length of documents in corpus $C$ (calculated in number of tokens), $\text{TF}(q_{i},D)$ is the term frequency of token $q_{i}$ within the set of tokens in $D$. The IDF is the inverse document frequency, which is calculated as:

$$\text{IDF}(q_{i},C)=\ln\left(\frac{|C|-\text{DF}(q_{i},C)+0.5}{\text{DF}(q_{i},C)+0.5}+1\right)$$

Where document frequency $\text{DF}(q_{i},C)$ is the number of documents in $C$ containing $q_{i}$. Although $B(Q,D)$ depends on the query, which is only given during retrieval, we show below how to reformulate the equation to eagerly calculate the TF and IDF during indexing.

Let’s now consider all tokens in a vocabulary $V$, denoted by $t\in V$. We can reformulate $S(t,D)$ as:

$$S(t,D)=\text{TF}(t,D)\cdot\text{IDF}(t,C)\frac{1}{\mathcal{D}}$$

When $t$ is a token that is not present in document $D$, then $\text{TF}(t,D)=0$, leading to $S(t,D)=0$ as well. This means that, for most tokens in vocabulary $V$, we can simply set the relevance score to 0, and only compute values for $t$ that are actually in the document $D$. This calculation can be done during the indexing process, thus avoiding the need to compute $S(q_{i},D)$ at query time, apart from straightforward summations.

Given our sparse matrix of shape $|V|\times|C|$, we can use the query tokens to select relevant rows, leaving us a matrix of shape $|Q|\times|C|$, which we can then sum across the column dimension, resulting in a single $|C|$-dimension vector (representing the score of the score of each document for the query).

We implement a sparse matrix in Compressed Sparse Column (CSC) format (scipy.sparse.csc_matrix)⁹⁹9https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csc_matrix.html, which provides an efficient conversion between the coordinate and CSC format. Since we slice and sum alongside the column dimension, this implementation is the optimal choice among sparse matrix implementations. In practice, we replicate the sparse operations directly using Numpy array.

To split the text, we use the same Regular Expression pattern used by Scikit-Learn Pedregosa et al. (2011) for their own tokenizers, which is r"(?u)\b\w\w+\b". This pattern conveniently parses words in UTF-8 (allowing coverage of various languages), with \b handling word boundaries. Then, if stemming is desired, we can stem all words in the vocabulary, which can be used to look up the stemmed version of each word in the collection. Finally, we build a dictionary mapping each unique (stemmed) word to an integer index, which we use to convert the tokens into their corresponding index, thus significantly reducing memory usage and allowing them to be used to slice Scipy matrices and Numpy arrays.

Upon computing scores for all documents in a collection, we can complete the search process by selecting the top-$k$ most relevant elements. A naive approach to this would be to sort the score vector and select the last $k$ elements; instead, we take the partition of the array, selecting only the last $k$ documents (unordered). Using an algorithm such as Quickselect Hoare (1961), we can accomplish this in an average time complexity of $O(n)$ for $n$ documents in the collection, whereas sorting requires $O(n\log n)$. If the user wishes to receive the top-$k$ results in order, sorting the partitioned documents would take an additional $O(k\log k)$, which is a negligible increase in time complexity assuming $k\ll n$. In practice, BM25S allows the use of two implementations: one based in numpy, which leverages np.argpartition, and another in jax, which relies on XLA’s top-k implementation. Numpy’s argpartition uses¹⁰¹⁰10https://numpy.org/doc/stable/reference/generated/numpy.argpartition.html the introspective selection algorithm Musser (1997), which modifies the quickselect algorithm to ensure that the worst-case performance remains in $O(n)$. Although this guarantees optimal time complexity, we observe that JAX’s implementation achieves better performance in practice.

We implement optional multi-threading capabilities through pooled executors¹¹¹¹11Using concurrent.futures.ThreadPoolExecutor to achieve further speed-up during retrieval.

Above, we describe how to implement BM25S for one variant of BM25 (namely, Lucene). However, we can easily extend the BM25S method to many variants of BM25; the sparsity can be directly applied to Robertson’s original design Robertson et al. (1995), ATIRE Trotman et al. (2014), and Lucene. For other models, a modification of the scoring described above is needed.

For benchmarking, we use the publicly available datasets from the BEIR benchmark Thakur et al. (2021). Results in Table 1 show that BM25S is substantially faster than Rank-BM25, as it achieves over 100x higher throughput in 10 out of the 14 datasets; in one instance, it achieves a 500x speedup. Further details can be found in Appendix A.

We further examine the impact of tokenization on each model in Table 2 by comparing BM25S Lucene with $k_{1}=1.5$ and $b=0.75$ (1) without stemming, (2) without stop words, and (3) with neither, and (4) with both. On average, adding a Stemmer improves the score on average, wheareas the stopwords have minimal impact. However, on individual cases, the stopwords can have a bigger impact, such as in the case of Trec-COVID (TC) and ArguAna (AG).

In Table 3, we compare many implementation variants, including commercial (Elasticsearch) offerings and reproducibility toolkits (Pyserini). We notice that most implementations achieve an average be between 39.7 and 40, with the exception of Elastic which achieves a marginally higher score. The variance can be attributed to the difference in the tokenization scheme; notably, the subword tokenizer used in BM25-PT likely lead to the difference in the results, considering the implementation is a hybrid between ATIRE and Lucene, both of which achieve better results with a word-level tokenizer. Moreover, although Elasticsearch is built on top of Lucene, it remains an independent commercial product, and the documentations¹⁴¹⁴14https://www.elastic.co/guide/en/elasticsearch/reference/current/index.html do not clearly describe how they are splitting the text¹⁵¹⁵15https://www.elastic.co/guide/en/elasticsearch/reference/current/split-processor.html, and whether they incorporate additional processing beyond the access to a Snowball stemmer and the removal of stopwords.