SparTerm: Learning Term-based Sparse Representation for Fast Text Retrieval

Bag-of-words(BoW) methods have played a central role in the first-stage retrieval. These methods convert a document or query into a set of single terms, and each term associates a weight to characterize its weight. Most of the early common practice adopted TF-IDF style models to calculate weights. Robertson (Robertson and Walker, 1994) proposed the well-known method BM25, which further improve the performance of the original TF-IDF. Later proposed methods, such as (Lafferty and Zhai, 2001), (Zaragoza et al., 2004), (Strohman et al., 2005), did not show much advantage over BM25. More recently, Hamed Zamani (Zamani et al., 2018) proposed SRNM to learn a sparse coding in hidden space using weak supervision, which shows good potential for solving the “lexical mismatch” problem. However, the latent unexplainable tokens can not ensure that documents with exact matched terms can be retrieved.

The pre-trained language models like BERT (Devlin et al., 2018) show new possibilities for text retrieval. Based on dense representations, Lee (Lee et al., 2019) proposed ORQA with bi-encoder architecture to retrieve candidate passages for question answering using FAISS (Johnson et al., 2017). However, analysis from (Luan et al., 2020) concludes that bi-encoders based on dense representation suffer from its capacity limitation in scenarios that emphasize long document retrieval and exact matching. Following the late-interaction paradigm, Khattab (Khattab and Zaharia, 2020) proposed Col-BERT to conduct efficient interaction between the query and document, which can run 150x faster than fully-interactive BERT but achieve comparable precision. Though much faster than BERT, Col-BERT is still not computationally feasible for large scale first-stage retrieval, for the existence of the late interaction layer.

Several PLM-based models have emerged to improve the traditional sparse BoW representations. Dai (Dai and Callan, 2019) proposed DeepCT to estimate a term’s weight considering its contextualized information, and this work was later extended to generate document-level term weights (Dai and Callan, 2020). Another work Doc2query (Nogueira et al., 2019) tries to “translate” potential queries to expand document content, which also shows a large improvement compared to the traditional BM25 method. The biggest difference between our work and these two methods is that DeepCT and Doc2Query train an auxiliary intermediate model to help refine the sparse representations, while SparTerm is desinged to directly learn sparse representations within the whole vocabulary.

Figure 2(a) depicts the general architecture of SparTerm which comprises an importance predictor and a gating controller. Given the original textual passage $p$, we aim to map it into a deep and contextualized sparse representation $p{{}^{\prime}}$ in the vocabulary space. The mapping process can be formulated as:

where $\mathcal{F}$ is the item importance predictor and $\mathcal{G}$ the gating controller. The importance predictor $\mathcal{F}$ generates a dense vector representing the semantic importance of each item in the vocabulary. The gating controller $\mathcal{G}$ generates a binary gating vector to control which terms to appear in the final sparse representation. To achieve this, we let $||\mathcal{G}(p)||<\lambda$ and $\mathcal{G}(p)\in\{0,1\}^{v}$, where $\lambda$ is the maximum number of non-zero elements for $p{{}^{\prime}}$, and $v$ the vocabulary size. These two modules cooperatively ensure the sparsity and flexibility of the final representation $p{{}^{\prime}}$. We discuss the detailed model architecture and learning strategy for $\mathcal{F}$ and $\mathcal{G}$ in the following sections.

Given the input passage $p$, the importance predictor outputs semantic importance of all the terms in the vocabulary, which unify term weighting and expansion into the framework. As shown in Figure  2(b), prior to importance prediction, BERT-based encoder is employed to help get the deep contextualized embedding $h_{i}$ for each term $w_{i}$ in the passage $p$. Each $h_{i}$ models the surrounding context from a certain position $i$, thus providing a different view of which terms are semantically related to the topic of the current passage. With a token-wise importance predictor, we obtain a dense importance distribution $I_{i}$ of dimension $v$ for each $h_{i}$:

where $Transform$ denotes a linear transformation with GELU activation and layer normalization, $E$ is the shared word embedding matrix and $b$ the bias term. Note that the token-wise importance prediction module is similar to the masked language prediction layer in BERT, thus we can initialize this part of parameters directly from pre-trained BERT. The final passage-wise importance distribution can be fetched simply by the summation of all token-wise importance distributions:

where $L$ is the sequence length of passage $p$ and Relu activation function is leveraged to ensure the nonnegativity of importance logits.

The gating controller generates a binary gating signal of which terms to activate to represent the passage. First, the terms appearing in the original passage, which we referred to as literal terms, should be activated by the controller by default. Apart from the literal terms, some other terms related to the passage topic are also expected to be activated to tackle the “lexical mismatch” problem of BOW representation. Accordingly, we propose two kinds of gating controller: literal-only gating and expansion-enhanced gating, which can be applied in scenarios with different requirements for lexical matching.

Literal-only Gating. If simply setting $\mathcal{G}(p)=BOW(p)$, where $BoW(p)$ denotes the binary BoW vector for passage $p$, we get the literal-only gating controller. In this setting, only those terms existing in the original passage are considered activated for the passage representation. Without expansion for non-literal terms, the sparse representation learning is reduced to a pure term re-weighting scheme. Nevertheless, in the experiment part, we empirically show that this gating controller can achieve competitive retrieval performance by learning importance for literal terms.

Exapnsion-enhanced Gating. The expansion-enhanced gating controller activates terms that can hopefully bridge the “lexical mismatch” gap. Similar to the importance prediction process formulated by Equation (2) and Equation (3), we obtain a passage-wise dense term gating distribution $G$ of dimension $v$ with independent network parameters, as shown in Figure 2(c). Note that although the gating distribution $G$ and the importance distribution $I$ share the same dimension $v$, they are different in logit scales and mathematical implications. $I$ represents the semantic importance of each term in vocabulary, while $G$ quantifies the probability of each term to participate in the final sparse representation. To ensure the sparsity of $p^{{}^{\prime}}$, we apply a binarizer to $G$:

where the $Binarizer$ denotes a binary activation function which outputs only 0 or 1. The gating vector for expansion terms $G_{e}$ is obtained by:

where the bitwise negation vector $\neg BoW(p)$ is applied to ensure orthogonal to the literal-only gating. Simply adding the expansion gating and the literal-only gating, we get the final expansion-enhanced gating vector $G_{le}$:

Involving both literal and expansion terms, the final sparse representation can be a “free” distribution in the vocabulary space. Note that in the framework of SparTerm, expanded terms are not directly appended to the original passage, but are used to control the gating signal of whether allowing a term participating the final representation. This ensures the input text to the BERT encoder is always the natural language of the original passage.

In this section, we introduce the training strategy of the importance predictor and expansion-enhanced gating controller.

Training the importance predictor. The importance predictor is trained end-to-end by optimizing the ranking objective. Let $R=\{(q_{1},p_{1,+},p_{1,-}),...,(q_{N},p_{N,+},p_{N,-})\}$ denote a set of N training instances; each containing a query $q_{i}$, a posotive candidate passage $p_{i,+}$ and a negative one $p_{i,-}$, indicating that $p_{i,+}$ is more relevant to the query than $p_{i,-}$. The loss function is the negative log likelihood of the positive passage:

where $q_{i}^{{}^{\prime}}$, $p_{i,+}^{{}^{\prime}}$, $p_{i,-}^{{}^{\prime}}$ is the sparse representation of $q_{i}$, $p_{i,+}$, $p_{i,-}$ obtained by Equation (1), $sim$ denotes any similarity measurement such as dot-product. Different with the training objective of DeepCT LABEL:dai2019context, we don’t directly fit the statistical term importance distribution, but view the importance as intermediate variables that can be learned by distant supervisory signal for passage ranking. End-to-end learning can involve every terms in the optimization process, which can yield smoother importance distribution, but also of enough distinguishability.

Training the exapnsion-enhanced gating controller. We summarize four types of term expansion in Table 1, all of which can be optimized in our SparTerm framework. Intuitively, the pre-trained BERT already has the ability of expanding synonym words and co-occured words by the Masked Language Model(MLM) pre-training task. Therefore, in this paper, we focus on expanding passage2query-alike and summarization terms. Given a passage-query/summary parallel corpus $\mathbf{C}$, where $p$ is a passage, $t$ the corresponding target text, and $T$ of dimension $v$ is the binary bag-of-words vector of $t$. We use the binary cross-entropy loss to maximize probability values of all the terms in vocabulary:

where $G$ is the dense gating probability distribution for $p$, $\lambda_{1}$ and $\lambda_{2}$ two tunable hyper-parameters. $\lambda_{1}$ is the loss weight for terms expected not to be expanded, while $\lambda_{2}$ is for terms that appear in the target text. In the experiment, we set $\lambda_{2}$ much larger than $\lambda_{1}$ to encourage more terms to be expanded.

End-to-end joint training. Intuitively, the supervisory ranking signal can also be leveraged to guide the training of the gating controller, thus we can train the importance predictor and gating controller jointly:

We evaluate our method on MSMARCO (Nguyen et al., 2016) which consists of two benchmark datasets:

MSMARCO Passage Retrieval dataset is based on the public MSMARCO dataset with a collection of 8.8M passages from Web pages gathered from Bing’s results to 1M real-world queries. Each query is associated with one or very few passages marked as relevant while no passage explicitly indicated as irrelevant. We build a small dev set for validating the full ranking performance instead of re-ranking by sampling the most relevant 1M passages to 1000 queries from the original passage ranking dev set with BM25.

MSMARCO Document Retrieval dataset is based on the source documents which contain the passages in the passage retrieval task. The dataset contains 367,013 documents and 367,013 queries for training set and 5,193 queries for dev set.

The original Dev Set of MSMARCO dataset is a re-ranking task, which is inconsistent with the retrieval task. Therefore, to find the best checkpoint of our model more accurately we build a new Dev Set to evaluate the retrieval performance by sampling about 1M passages from the collections and 1,000 queries from the original Dev Set (including the top 1000 passages of each query retrieved by BM25). To evaluate the full ranking performance of our model, we use the sparse representation of each document to build the inverted index and use the sparse representation of queries to retrieval topK relevant documents and measure the performance with MRR@10 and Recall from top10 to top1000.

The Importance Predictor and Gating Controller of our model have the same architecture and hyper-parameters of BERT (12-layer, 768-hidden, 12-heads, 110M parameters) and do not share weights. We initialize the Importance Predictor with Google’s official pre-trained $\mathrm{BERT_{base}}$ model while the parameters of Token-wise Importance Predictor are initialized with the Masked Language Prediction layer of BERT. When using expansion-enhanced gating, the Gating Controller is also initialized with $\mathrm{BERT_{base}}$. We fine-tune our model on the training set of MSMARCO passage retrieval dataset on 4 NVIDIA-v100 GPUs with a batch size of 128. During the fine-tuning, we first fine-tune the Gating Controller with Equation (8) for 50k iterations where $\lambda_{1}=1e-3$ and $\lambda_{2}=1$. Then we fix the parameters of the Gating Controller and fine-tune our SparTerm jointly for 100k iterations. We use Adam optimizer with the learning rate $2\times 10^{-5}$. To ensure the sparsity, the threshold in the Binarizer in Equation (4) is set to $0.7$. We do not fine-tune our model on the training set of document retrieval dataset but just use the model trained on the passage retrieval dataset for the document ranking.

We compare our model with the following strong baselines which are all methods based on sparse representation . The former two focus on re-weighting while the latter two focus on document expansion:

• •

  BM25(Robertson and Walker, 1994) is a bag-of-words retrieval models with frequency-based signals to estimate the weights of terms in a text.

• •

  DeepCT(Dai and Callan, 2019) is a deep contextualized term weighting model which maps the BERT’s representations to term weightings for retrieval.

• •

  Doc2query(Nogueira et al., 2019) is a document expansion method with Transformer that can expand documents with terms related to the documents’ content.

• •

  Doc2query-T5(Cheriton, 2019) is a document expansion method which utilizes more powerful T5 (Raffel et al., 2019) language model to generate queries for document expansion.

We also evaluate three different settings of SparTerm for evaluation:

• •

  SparTerm(literal-only) uses Importance Predictor with the Literal-only Gating which can also be seen as a term weighting model.

• •

  SparTerm(expansion-only) uses the Expansion-enhanced Gating for passage expansion without term weighting. We just add the expanded words (weight of each word is 1) to the passages.

• •

  SparTerm(expansion-enhanced) implements both Importance Predictor and Expansion-enhanced Gating for sparse representation of passage.

Table 2 shows the full ranking performances of our models and baselines on MSMARCO Passage Retrieval dataset. SparTerm (expansion-enhanced) outperforms all baselines on MRR, achieving the state-of-the-art ranking performance among all sparse models, and outperforms all baselines except Doc2query-T5 on Recall. We find that SparTerm achieves more significant performance improvements on MRR and Recall@10-100, which illustrates that our model has a more significant ability on top ranking compared with previous sparse models. Further, pre-trained language model(PLM) based methods (DeepCT, Doc2query-T5, and SparTerm) perform better than those without PLM, demonstrating that PLM can facilitate the passage full ranking with better representation. Considering the improvements T5 brings to Doc2query, we believe that SparTerm can be further improved with more advanced PLM.

Even without any expansion, SparTerm(literal-only) outperforms DeepCT on both MRR and Recall, demonstrating that SparTerm can produce more effective term weights thus facilitating the retrieval. We also analyze the difference between SparTerm and DeepCT on term weighting in Section 5.4. With only the expanded words, SparTerm achieves a definite improvement compared with BM25, especially on Recall. This improvement proves the effectiveness of passage expansion on improving the Recall for retrieval.

For the Document Ranking task, we cut down each document into several passages to adapt the max length (256) of the sequence of our model and generate the sparse representation of each passage with our model. We compare our models with two baseline methods: BM25 (Robertson and Walker, 1994) and HDCT (Dai and Callan, 2020). HDCT is based on the work of DeepCT and focuses on document ranking, which is also a term weighting method. HDCT compares two different ways to combine the representations of passages for document ranking. The first one represents the document as a sum of the passage representations while the second one uses a decayed weighted sum. The PassageRetrievalMax does not represent the document but just calculates the scores of passages in the document and choose the maximum score as the score of the document for ranking. Table 3 shows the ranking performance of baselines and our models. Here we only report the results of PassageRetrievalMax of our models.

Strictly speaking, it is incomparable between HDCT and our models since we fine-tune SparTerm on MSMARCO pasage ranking dataset while HDCT was trained using document titles on MARCO. Even though, SparTerm(expansion-enhanced) still achieves a better performance on document ranking compared with HDCT, demonstrating that the sparse representation produced by SparTerm can also facilitate long document retrieval.

We conduct experiments to evaluate the performance of SparTerm with different query representation methods:

• •

  Query-tf is a one-tower model that use tf-based vectors to represent the queries while use the model to represent documents.

• •

  Query-neural-symmetric is a symmetric two-tower model to represent queries and passages that the two towers with the same architectures share the same weights.

• •

  Query-neural-asymmetric is a asymmetric two-tower model that the two towers do not share weights. Queries and passages are represented with different towers.

The results are reported in Table 4, from where we find that the neural representation of queries with the symmetric two-tower model brings better performance on MRR and Recall on our built Dev set. The symmetric model performs better than the asymmetric model might because asymmetric two-tower architecture leads to twice the quantity of parameters, which makes the model more difficult to converge. We further analyze the distribution of passage term weights with different query representation methods and find that tf-based representation of query results in a sharper distribution compared to the neural representation. The reason may be that the query representation is fixed during training, the model needs to give more weights to the relative terms in the positive passage.

To further evaluate the ability of SparTerm on term weighting, we normalize the term weights of passages weighted by DeepCT and SparTerm(literal-only) to the same range and visulization them in Figure 3. Figure 3 shows three different queries(the first column) and the most relevant passages. The depth of the color represents the weights of terms, deeper is higher. We find that both DeepCT and SparTerm can figure out the most important terms and give them higher weights. However, DeepCT obtains sparser and sharper distributions and only activates very few terms in a passage, missing some important terms, such as “allergic reaction” in the first case. SparTerm can yield a smoother importance distribution by activating more terms though not appearing in the query. This distribution allows the passage to be retrieved by more queries. This also demonstrates that our model has a better ability on pointing out important terms in a passage.

Figure 3 shows the expanded terms and their probabilities for different passages predicted by the Gating Controller. The probability of each term illustrates how likely this term to be expanded. It is obvious that our model can really activate some important terms not appearing in the passage but very semantically similar, especially occurring in the queries such as “sign” in the first case and “temperature” in the second case.

In order to analyze how these words are expanded and which category in Figure 3 do they belong to, we trace the source of each expanded word and show the top 5 words with their logits which contribute to the expanded word in Figure 4. We can find that there are basically three different situations of the expanded terms:

(1) The passage2query terms such as “temperature”: Almost every word in the passage contributes much to this kind of terms, which seem more likely to learn from the supervised signal.

(2) Synonyms of the original terms, i.e. “weather” and “climate”, “rainfall” and “rain”, “season, monthly” and “month”, “heat” and “hot”.

(3) Co-occurred words for the original terms, i.e. “season, heat”-¿“summer”, “wet, humidity, weather”-¿“rain” and “heat, rainfall, humidity”-¿“tropical, monsoon”.

The first situation is benefited by the optimization objective of the Gating Controller while the latter two are more likely the ability of MLM pretraining task since we reuse the MLM module for prediction in the Gating Controller.