SPLADE: Sparse Lexical and Expansion Model for First Stage Ranking

Introduction

The release of large pre-trained language models like BERT (Devlin et al. 2018) has shaken-up Natural Language Processing and Information Retrieval. These models have shown a strong ability to adapt to various tasks by simple fine-tuning. At the beginning of 2019, Nogueira and Cho (Nogueira and Cho 2019) achieved state-of-the-art results – by a large margin – on the MS MARCO passage re-ranking task, paving the way for LM-based neural ranking models. Because of strict efficiency requirements, these models have initially been used as re-rankers in a two-stage ranking pipeline, where first-stage retrieval – or candidate generation – is conducted with bag-of-words models (e.g. BM25) that rely on inverted indexes. While BOW models remain strong baselines (Yang et al. 2019), they suffer from the long standing vocabulary mismatch problem, where relevant documents might not contain terms that appear in the query. Thus, there have been attempts to substitute standard BOW approaches by learned (neural) rankers. Designing such models poses several challenges regarding efficiency and scalability: therefore there is a need for methods where most of the computation can be done offline and online inference is fast. Dense retrieval with approximate nearest neighbors search has shown impressive results (Xiong et al. 2021; Lin et al. 2020; Ding et al. 2020), but is still combined with BOW models because of its inability to explicitly model term matching. Hence, there has recently been a growing interest in learning sparse representations for queries and documents (Zamani et al. 2018; Dai and Callan 2019; Nogueira et al. 2019; Zhao et al. 2020; Bai et al. 2020). By doing so, models can inherit from the desirable properties of BOW models like exact-match of (possibly latent) terms, efficiency of inverted indexes and interpretability. Additionally, by modeling implicit or explicit (latent, contextualized) expansion mechanisms – similarly to standard expansion models in IR – these models can reduce the vocabulary mismatch.

The contributions of this paper are threefold: (1) we build upon SparTerm (Bai et al. 2020), and show that a mild tuning of hyperparameters brings improvements that largely outperform the results reported in the original paper; (2) we propose the SParse Lexical AnD Expansion (SPLADE) model, based on a logarithmic activation and sparse regularization. SPLADE performs an efficient document expansion (Bai et al. 2020; MacAvaney et al. 2020), with competitive results with respect to complex training pipelines for dense models like ANCE (Xiong et al. 2021); (3) finally, we show how the sparsity regularization can be controlled to influence the trade-off between efficiency (in terms of the number of floating-point operations) and effectiveness.

Related Works

Dense retrieval based on BERT Siamese models (Reimers and Gurevych 2019) has become the standard approach for candidate generation in Question Answering and IR (Guu et al. 2020; Karpukhin et al. 2020; Xiong et al. 2020; Lin et al. 2020; Ding et al. 2020). While the backbone of these models remains the same, recent works highlight the critical aspects of the training strategy to obtain state-of-the-art results, ranging from improved negative sampling (Xiong et al. 2020; Ding et al. 2020) to distillation (Hofstätter et al. 2020; Lin et al. 2020). ColBERT (Khattab and Zaharia 2020) pushes things further: the postponed token-level interactions allow to efficiently apply the model for first-stage retrieval, benefiting of the effectiveness of modeling fine-grained interactions, at the cost of storing embeddings for each (sub)term – raising concerns about the actual scalability of the approach for large collections. To the best of our knowledge, very few studies have discussed the impact of using approximate nearest neighbors (ANN) search on IR metrics (Boytsov 2018; Tu et al. 2020). Due to the moderate size of the MS MARCO collection, results are usually reported with an exact, brute-force search, therefore giving no indication on the effective computing cost.

Sparse Lexical representations for first-stage ranking

In this section, we first describe in details the SparTerm model (Bai et al. 2020), before presenting our model named SPLADE.

SparTerm predicts term importance – in BERT WordPiece vocabulary ($|V|=30522$) – based on the logits of the Masked Language Model (MLM) layer. More precisely, let us consider an input query or document sequence (after WordPiece tokenization) $t=(t_{1},t_{2},...,t_{N})$, and its corresponding BERT embeddings $(h_{1},h_{2},...,h_{N})$. We consider the importance $w_{ij}$ of the token $j$ (vocabulary) for a token $i$ (of the input sequence):

where $E_{j}$ denotes the BERT input embedding for token $j$, $b_{j}$ is a token-level bias, and transform$(.)$ is a linear layer with GeLU activation and LayerNorm. Note that Eq. 1 is equivalent to the MLM prediction, thus it can be also be initialized from a pre-trained MLM model. The final representation is then obtained by summing importance predictors over the input sequence tokens, after applying ReLU to ensure the positivity of term weights:

where $g_{j}$ is a binary mask (gating) described latter. The above equation can be seen as a form of query/document expansion, as observed in (Bai et al. 2020; MacAvaney et al. 2020), since for each token of the vocabulary the model predicts a new weight $w_{j}$. SparTerm (Bai et al. 2020) introduces two sparsification schemes that turn off a large amount of dimensions in query and document representations, allowing to efficiently retrieve from an inverted index:

lexical-only is a BOW masking, i.e. $g_{j}=1$ if token $j$ appears in $t$, and 0 otherwise;

expansion-aware is a lexical/expansion-aware binary gating mechanism, where $g_{j}$ is learned. To preserve the original input, it is forced to 1 if the token $j$ appears in $t$.

Let $s(q,d)$ denote the ranking score obtained via dot product between $q$ and $d$ representations from Eq. (2). Given a query $q_{i}$, a positive document $d_{i}^{+}$ and a negative document $d_{i}^{-}$, SparTerm is trained by minimzing the following loss:

SparTerm expansion-aware gating is somewhat intricate, and the model cannot be trained end-to-end: the gating mechanism is learned beforehand, and fixed while fine-tuning the matching model with $\mathcal{L}_{rank}$, therefore preventing the model to learn the optimal sparsification strategy for the ranking task. Moreover, the two lexical and expansion-aware strategies do perform almost equally well, questioning the actual benefits of expansion.

2. SPLADE: SParse Lexical AnD Expansion model

In the following, we propose slight, but essential changes to the SparTerm model that dramatically improve its performance.

We introduce a minor change in the importance estimation from Eq. 2, by introducing a log-saturation effect which prevents some terms to dominate and naturally ensures sparsity in representations:

While it is intuitive that using a log-saturation prevents some terms from dominating – drawing a parallel with axiomatic approaches in IR and $\log$(tf) models (Fang et al. 2004) – the implied sparsity can seem surprising at first, but, according to our experiments, it obtains better experimental results and allows already to obtain sparse solutions without any regularization.

Ranking loss

Given a query $q_{i}$ in a batch, a positive document $d_{i}^{+}$, a (hard) negative document $d_{i}^{-}$ (e.g. coming from BM25 sampling), and a set of negative documents in the batch (positive documents from other queries) $\{d_{i,j}^{-}\}_{j}$, we consider the ranking loss from (Ding et al. 2020), which can be interpreted as the maximization of the probability of the document $d_{i}^{+}$ being relevant among the documents $d_{i}^{+},d_{i}^{-}$ and $\{d_{i,j}^{-}\}$:

The in-batch negatives (IBN) sampling strategy is widely used for training image retrieval models, and has shown to be effective in learning first-stage rankers (Karpukhin et al. 2020; Ding et al. 2020; Lin et al. 2020).

Learning sparse representations

Overall loss

We propose to combine the best of both worlds for end-to-end training of sparse, expansion-aware representations of documents and queries. Thus, we discard the binary gating in SparTerm, and instead learn our log-saturated model (Eq. 4) by jointly optimizing ranking and regularization losses:

Experimental setting and results

We trained and evaluated our models on the MS MARCO passage ranking dataset https://github.com/microsoft/MSMARCO-Passage-Ranking in the full ranking setting. The dataset contains approximately $8.8$M passages, and hundreds of thousands training queries with shallow annotation ($\approx 1.1$ relevant passages per query in average). The development set contains $6980$ queries with similar labels, while the TREC DL 2019 evaluation set provides fine-grained annotations from human assessors for a set of $43$ queries (Craswell et al. 2020).

We initialized the models with the BERT-base checkpoint. Models are trained with the ADAM optimizer, using a learning rate of $2e^{-5}$ with linear scheduling and a warmup of $6000$ steps, and a batch size of $124$. We keep the best checkpoint using MRR@10 on a validation set of $500$ queries, after training for $150$k iterations (note that this is not optimal, as we validate on a re-ranking task). We consider a maximum length of $256$ for input sequences. In order to mitigate the contribution of the regularizer at the early stages of training, we follow (Paria et al. 2020) and use a scheduler for $\lambda$, quadratically increasing $\lambda$ at each training iteration, until a given step ($50$k in our case), from which it remains constant. Typical values for $\lambda$ fall between $1e^{-1}$ and $1e^{-4}$. For storing the index, we use a custom implementation based on Python arrays, and we rely on Numba (Lam et al. 2015) to parallelize retrieval. Models We made the code public at https://github.com/naver/splade are trained using PyTorch (Paszke et al. 2019) and HuggingFace transformers (Wolf et al. 2020), on $4$ Tesla $V100$ GPUs with $32$GB memory.

Evaluation

Results are given in Table 1. Overall, we observe that: (1) our models outperform the other sparse retrieval methods by a large margin (except for recall@1000 on TREC DL); (2) the results are competitive with state-of-the-art dense retrieval methods.

Effectiveness-efficiency trade-off

The role of expansion

Conclusion

Recently, dense retrieval based on BERT has demonstrated its superiority for first-stage retrieval, questioning the competitiveness of traditional sparse models. In this work, we have proposed SPLADE, a sparse model revisiting query/document expansion. Our approach relies on in-batch negatives, logarithmic activation and FLOPS regularization to learn effective and efficient sparse representations. SPLADE is an appealing candidate for initial retrieval: it rivals the latest state-of-the-art dense retrieval models, its training procedure is straightforward, its sparsity/FLOPS can be controlled explicitly through the regularization, and it can operate on inverted indexes. In reason of its simplicity, SPLADE is a solid basis for further improvements in this line of research.

References