Bridging the Gap Between Indexing and Retrieval for Differentiable Search Index with Query Generation

The original DSI method exhibits a gap between the input data used at indexing and that used at retrieval. In order to bridge this gap and improve DSI’s effectiveness, we propose DSI-QG which uses a query generation model to generate a set of potentially-relevant queries to represent each candidate document for indexing. Specifically, we denote $\hat{Q_{i}}$ as the set of queries generated by a query generation model $QG$ given the document $d_{i}$:

All the generated queries $\hat{q}$ in $\hat{Q_{i}}$ share the same docid as $d_{i}$, and $|\hat{Q_{i}}|=n$. We then replace the original documents that need to be indexed with their corresponding generated queries, i.e. using $\hat{Q_{i}}$ in place of $d_{i}$. In other words, a document is replaced by the set of queries generated for that document. Thus, during the indexing phase in DSI-QG, the modified DSI model is trained to associate the generated queries of each candidate document with its docid:

The retrieval phase of DSI-QG is the same as the original DSI model and takes a user query as input and uses beam search to construct a ranked docid list. Note that each query in $\hat{Q_{i}}$, that was generated for $d_{i}$, is used separately for the other queries for $d_{i}$, i.e. queries for a document are not concatenated or combined into a single input. In summary, in our DSI-QG framework, a DSI model only observes short queries as input data during both indexing and retrieval thus eliminating the data distribution mismatch problem that affects the original DSI model.

A key factor for the success of the DSI-QG indexing framework is the query generation (QG) model. This model should generate high-quality and diverse relevant queries so that they can effectively represent the corresponding document from which they are generated. For this purpose, we train a T5 model with a similar seq2seq objective as Eq. (1), but in this case the input is the document text and the output is the labeled relevant query $q_{j}$:

After training a QG model, instead of using beam search, we use a random sampling strategy to generate a set of queries for each candidate document. This is because we find random sampling gives more creative and diverse queries than beam search, thus potentially covering more relevant information about the document. To avoid randomly generating irrelevant content and grammatically incorrect queries, we adopt the top-$k$ sampling scheme (Fan et al., 2018) which only allows the $k$ most likely next words to be sampled during the query generation and the probability mass of those $k$ next words is redistributed. In our experiments, we set $k=10$.

Intuitively, a document may be relevant to more than just one query, thus another factor that might impact the effectiveness of our proposed DSI-QG method is the number of generated queries $n$ to represent each document: we discuss the impact of $n$ in the result section.

To generalize our DSI-QG framework to the cross-lingual IR setting, we also train a multi-lingual T5 model (Xue et al., 2021) to generate queries in different languages, and then in turn use these to represent a document. To achieve this, we use a prompt-based approach to control the generated query language. Specifically, we place the target language and the document text in the following template for both training and inference:

where [lang] and [doc] is the placeholder for the target query language and the document text. In our cross-lingual experiments, [doc] is always written in English and [lang] is a language other than English. We generate multiple queries for all the target languages and use these to represent each English document. By doing so, our DSI-QG model can learn to build connections between the English documents identifiers with queries from different languages, thus allowing to perform cross-lingual retrieval with our proposed cross-lingual query generation model.

Although our query generation model adopts the top-$k$ sampling scheme to balance the relevance and diversity of generated queries, it still inevitably generates irrelevant queries due to the randomness of the sampling process. This problem is even more considerable when there is not enough training data to train the query generation model or when the model is ill trained. To further mitigate this problem, we add a cross-encoder ranker to rank all the generated queries and only use the top-$m$ ranked queries to represent the original document.

Specifically, we use monoBERT (Nogueira et al., 2019) as a cross-encoder ranker: this is a transformer encoder-based model that employs BERT and that takes a query-document pair (separated by a special [SEP] token) as input and outputs a relevance score $s$:

We train the ranker with supervised contrastive loss, similar to Gao et al. (Gao et al., 2021b):

In our cross-encoder ranker, all the query tokens can interact with all the document tokens thus it has more substantial relevance modeling power than other ranker architectures, such as dual- or bi-encoders (Lin et al., 2020; Ren et al., 2021; Zhang et al., 2022a; Lu et al., 2022). We then rank all the generated queries for each document in decreasing order of the relevance score estimated by our ranker. From this ranking, we only select the top-$m$ queries to pass to the downstream DSI indexing training, thus effectively filtering out the remaining $n-m$ queries. We note that our query generation model and cross-encoder ranker are large transformer models and thus require substantial computational resources in addition to the DSI model alone. However, these additional computations only happen during the offline indexing phase, and will not affect the online query latency. We leave methods for reducing the computational resources required for indexing to future work.

Following the original DSI paper, we conduct our experiments on subsets of publicly available document retrieval datasets, namely NQ 320k (Kwiatkowski et al., 2019), for the mono-lingual document retrieval task, and XOR QA 100k (Asai et al., 2021), for the cross-lingual retrieval task. The NQ 320k dataset has $\approx$307k training query-document pairs and $\approx$8k dev query-document pairs. All the queries and documents in NQ 320k are in English. We follow the description in DSI (Tay et al., 2022) and SEAL (Bevilacqua et al., 2022) to construct the dataset as the code for dataset construction is not yet publicly available at the time of writing. For XOR QA 100k, we use the gold paragraph data available in the original repository¹¹1https://github.com/AkariAsai/XORQA#gold-paragraph-data which contains around 15k gold (annotated as relevant) document-query pairs in the training set and 2k gold document-query pairs in the dev set. Queries in both train and dev sets are in 7 typologically diverse languages²²2They are Ar, Bn, Fi, Ja, Ko, Ru and Te. and documents are in English. The total number of documents in the XOR QA training set and dev set is around 17k. This is a very small number of documents, likely to render the retrieval task too easy. We then randomly sample 93k documents from a dump of the English Wikipedia corpus to form a 100k collection for testing our models, thus increasing how challenging retrieval in this collection is.

We compare DSI-QG with the following baselines:

• •

  BM25 (Robertson and Zaragoza, 2009): a classic sparse retrieval method based on inverted indexes. This method usually only works for mono-lingual retrieval tasks as it is a keyword-matching method. We use the Pyserini (Lin et al., 2021a) implementation of BM25 for this baseline.

• •

  BM25 + docT5query (Nogueira and Lin, 2019): a sparse retrieval method which also leverages query generation. It uses a T5 model to generate a set of queries and appends them to the original document. Then it uses an inverted index and BM25 to retrieve augmented documents. In the original study that investigated this method, only the mono-lingual retrieval task was considered (Nogueira and Lin, 2019). For fair comparison with DSI-QG, we adapt this method to the cross-lingual retrieval setting by replacing the mono-lingual T5 query generation model with the same multi-lingual T5 generation model used in our DSI-QG. We also use the Pyserini implementation for this baseline.

• •

  SEAL (Bevilacqua et al., 2022): an autoregressive generation model that is similar to DSI. It treats ngrams that appear in the collection as document identifiers; at retrieval time, it directly generates and scores distinctive ngrams that are mapped to the documents. Unlike DSI, which unifies the index into the model parameters, SEAL requires a separate index data structure to perform an efficient search. Note that no publicly available implementation of SEAL currently exists. Unlike for DSI below, the re-implementation of SEAL is outside the scope of our work, and thus we report the results obtained by Bevilacqua et al. (Bevilacqua et al., 2022) on the NQ 320k dataset. SEAL has not been devised for and experimented with the task of cross-lingual retrieval and thus no results for XOR QA 100k are reported.

• •

  mDPR (Karpukhin et al., 2020; Asai et al., 2021): a mBERT-based cross-lingual dense passage retrieval method trained with a contrastive loss and with hard negatives sampled from the top passages retrieved by BM25. mDPR relays on nearest neighbor index search (Faiss implementation (Johnson et al., 2019)) to retrieve the passages that have the closest embeddings to the query embedding. We train the mDPR model with the Tevatron dense retriever training toolkit (Gao et al., 2022). Of course, due to its cross-lingual nature, we run mDPR only on the cross-lingual dataset, XOR QA 100k.

• •

  DSI (Tay et al., 2022): The original DSI method that uses documents text as input for indexing. Since the original code has not currently been made available by the authors, we implement and train the DSI model ourselves using the Huggingface transformers Python Library. We provide the implementation of this DSI model in our public code repository, along with the implementations of the other models considered in this paper.

Following the original DSI paper, for both datasets, we evaluate baselines and our models on the dev set with Hits@1 and Hits@10. This metric reports the proportion of the correct docids ranked in the top 1 and top 10 predictions. In addition, for XOR QA 100k we also report nDCG@10; this metric is not available for NQ 320k for some of the considered baselines and thus we do not report it as comparisons between methods cannot then be made.

There are three Transformer models in our DSI-QG framework: a query generation model, a cross-encoder ranker, and a DSI model.

For the NQ 320k dataset, we fine-tune an existing docT5query query generation model checkpoint³³3https://huggingface.co/castorini/doc2query-t5-large-msmarco with the training portion of the NQ 320k dataset. For the cross-encoder ranker, we train a ‘BERT-large-uncased’ checkpoint and 15 hard negatives documents sampled from BM25. For the DSI model, we use the standard pre-trained T5 model (Raffel et al., 2019) to initialize the model parameters.

For XOR QA 100k, we use the multi-lingual T5 model (Xue et al., 2021) to initialize both the query generation model and DSI model. For the cross-lingual ranker, we train ‘xlm-roberta-large’ (Conneau et al., 2020) checkpoint with BM25 hard negatives provided by the XOR QA official repository. For our trained query generation model, we train the model with a batch size of 128 and a learning rate of $1e^{-4}$ with Adam optimizer for 600 training steps on XOR QA 100k datasets,which is equivalent to about 6 epochs, and 9500 steps on the NQ 320k dataset which is equivalent to about 4 epochs. The DSI models in our DSI-QG method are trained for a maximum of 1M steps with a batch size of 256 and a learning rate of $5e^{-5}$ with 100k warmup-steps. Since the documents in DSI-QG are represented by generated short queries, we set the maximum length of the input data to 32 tokens for faster training and saving GPU memory usage. For training the original DSI model, we use the training configuration suggested in the original paper (Tay et al., 2022). For mDPR trained on XOR QA, we follow the training configuration in the XOR QA paper (Asai et al., 2021), which uses a multi-lingual BERT-base model as the backbone query and passage encoder. All Transformer models used in this paper are implemented with Huggingface transformers (Wolf et al., 2020) and training is conducted on 8 Tesla A100 GPUs.

We start by discussing the effectiveness of the proposed DSI-QG framework on the mono-lingual retrieval task; recall that these experiments are based on the NQ 320k English mono-lingual retrieval dataset.

Table 4.3 contains the Hits scores of the baselines and our DSI-QG methods on NQ 320k. For BM25 + docTquery and DSI-QG, we first generated $n=100$ queries for each document; then we ranked them using the cross-encoder ranker and select only the top $m=50$ queries. This process thus resulted in 50 queries being used to represent each document. To explore the impact of different model sizes, we report the results for DSI and DSI-QG with T5-base (200M parameters) and T5-large (800M).

The results show that the original DSI method performs worse than other baselines, with the exception of DSI with T5 large which outperforms BM25 on both Hits scores and SEAL on Hits@1. BM25 with docT5query document augmentation, which is a simple and straightforward way of leveraging query generation, achieves the best Hits@1 among the baselines we consider. These results suggest that the existing autoregressive generation-based information retrieval methods are inferior to the considered baselines in the mono-lingual retrieval task.

On the other hand, our DSI-QG outperforms all baselines by a large margin on both Hits measures. Compared to the original DSI method, Hits@1 and Hits@10 improve by 132% and 46% for T5-base, and 83% and 32% for T5-large. This suggests that the query generation employed in DSI-QG successfully addresses the data distribution mismatch problem that afflicts the original DSI method.

We further compare our results with those reported by Wang et al. ([n.d.]), who developed the DSI-based technique in parallel with ours. On NQ320k they report their query-generation methods based on a docT5query fine-tuned on the dataset achieves Hits@10 values of 88.48 and 88.45 for the base and large backbones, respectively. This result supports the superiority of DSI models that exploit query generation, and confirms the backbone model size has minor influence on effectiveness.

Next, we specifically focus on the impact of model size on retrieval effectiveness. We note that the effectiveness of the original DSI method decreases dramatically with a smaller base model. In contrast, model size has relatively little impact on the effectiveness of DSI-QG. This suggests that when using the DSI-QG framework, a large pre-trained T5 model is not necessarily required. The use of a smaller T5 model means that DSI-QG can feature faster retrieval inference time and lower GPU memory requirements.

Next we examine the effectiveness of the proposed DSI-QG framework on the cross-lingual retrieval task; recall that these experiments are based on the XOR QA 100k cross-lingual dataset.

In Table 2(c), we report the results obtained across the different languages. For BM25 + docTquery and DSI-QG, we first generated $n=700$ queries for each document (100 per language) and then separately rank the generated queries for each language using the cross-encoder ranker and the cutoff $m=10$.This resulted in 70 generated queries being used to represent each document (10 for each language).

The results show that the original DSI model performs much worse on XOR QA 100k than on NQ 320k (cfr. Table 4.3). In fact, across many languages, DSI-base fails to retrieve any relevant document in the top rank position (Hits@1). This is likely due to the data distribution mismatch problem being further exacerbated by the language gap in the cross-lingual document retrieval task (Zhang et al., 2022b). In contrast, our proposed DSI-QG achieves the highest Hits values across all languages with the only exceptions that its Hit@10 on Finnish is lower than that of mDPR, as are the nDCG@10 values for Arab and Russian.

These results suggest that, with a cross-lingual query generation model, our DSI-QG not only can address the indexing and retrieval gap brought by the data type mismatch but can also address the gap brought by the language mismatch that instead affects the original DSI model.

Next, we discuss the effect of different components and factors on our DSI-QG model. Specifically, we study the effect of the cross-encoder ranker and the impact of the rank cut-off $m$ used when ranking and selecting the generated queries.

DSI-QG involves a step of query generation and further ranking and selection of queries to represent a document at indexing. In Table 3 we report an example of a document from the XOR QA dataset, one of the multilingual query for which this document has been assessed as relevant in the collection, and a sample of the queries that are generated by DSI-QG for the same target language (Russian). The sample of the queries are ordered according to the scores generated for these queries by the cross-encoder ranker. While all the top 3 queries would be used by DSI-QG to represent the document at indexing (when $m>5$), the bottom 3 queries would be discarded by all DSI-QG settings we experimented with in the paper, except when $m=n=100$.

We then generalise the above analysis by considering all queries that have been generated for all documents in the dataset. Figure 4 shows the effectiveness, measured in terms of mean reciprocal rank (MRR) of each of the generated queries (ordered by the cross-encoder ranker) at retrieving the relevant document when retrieval is performed using the mDPR baseline. Recall that mDPR is generally highly effective on this dataset, as seen in Table 2(c)(c). Thus, we use mDPR to provide an estimation of “query quality” as in this way we decouple this estimation from the training of DSI-QG. The trend observed in the plot suggests that the quality of the generated queries decreases as their rank assigned by the cross-encoder ranker increases, i.e. generally queries in early rank positions are associated to higher mDPR effectiveness than queries at later rank positions, attesting to the importance of ranking queries for use in DSI-QG.

While the original DSI uses the full length of a document as input to the Transformer used for indexing the document, DSI-QG uses $m$ queries to represent a document, each of them passed separately as input to the DSI Transformer. We argued that the effectiveness of the original DSI method is limited by the mismatch between the length of the input provided at indexing (documents, usually long) and retrieval (queries, usually short). The new framework we devised, DSI-QG, uses queries for indexing in lieu of documents: this aligns the lengths of the input at indexing (now generated queries, usually short) and the input at retrieval (real queries, usually short).

We then analyze the input lengths of DSI and DSI-QG to demonstrate that indeed DSI-QG’s indexing lengths are shorter and more aligned with the query lengths observed at retrieval. Input lengths are measured according to the T5 model tokenizer used in DSI, i.e. the number of tokens T5 produces for a text input. These statistics are reported in Table 4, and show that indeed for DSI input lengths greatly differ at indexing and retrieval, while these are similar in DSI-QG.

Pretrained transformer-based autoregressive generation models have been shown effective across many NLP tasks (Raffel et al., 2019; Brown et al., 2020). Recent studies also explored adapting this type of models to the information retrieval task.

Cao et al. have applied autoregressive generation models to conduct entity retrieval where queries are texts with a mention span and documents are entities in a knowledge base (Cao et al., 2021; Cao et al., 2022). In this setting, the documents’ identifiers are the entity names.

Different from the entity retrieval setting, Tay et al. proposed the differentiable search index (DSI) scheme (Tay et al., 2022), which is an autoregressive generation model trained to perform ad hoc document retrieval where the input of the model is a natural language query and the model regressively generates documents’ identifier strings that are potentially relevant to the given query.

In another direction,  Bevilacqua et al. proposed the SEAL model which treats ngrams that appear in the collection as document identifiers (Bevilacqua et al., 2022). At inference time, SEAL directly generates ngrams which can be used to score and retrieve documents stored in an FM-index (Ferragina and Manzini, 2000).

In contrast to the original DSI scheme and SEAL, our work focuses on augmenting document representations at indexing time so to bridge the gap between indexing and retrieval in the existing autoregressive generation IR systems.

The method of Wang et al. ([n.d.]), which was developed in parallel with ours and with no communication between the research groups, shares some commonalities with our DSI-QG in that it also uses query generation for augmenting document representations. However, while their work only considers English document retrieval and thus monolingual augmentation, our work goes beyond that and delves into the setting of cross-lingual document retrieval, offering a more comprehensive approach to the field.

Our DSI-QG indexing framework relies on a cross-encoder model to rank all generated queries in order to identify high-quality queries to represent documents. The intuition behind this design is that, for deep learning models, the generation task is usually a harder task than the classification task. Thus, many deep generation models follow the generate-then-rank paradigm to improve the generation quality. For example, the recent text-to-image generation model DALL$\cdot$E (Ramesh et al., 2021) also uses a ranker called CLIP (Radford et al., 2021) to rank all generated images and only present to the users the top-ranked images. On the other hand, while the GPT-3 language model (Brown et al., 2020) has been shown to perform poorly in solving mathematical problems (Hendrycks et al., 2021),  Cobbe et al. (Cobbe et al., 2021) have found that simply training verifiers to judge the correctness of the solutions generated by GPT-3 can significantly improve the success of the model for this task. DSI-QG can be seen as following this generate-then-rank paradigm.

Our DSI-QG framework relies on the query generation model to generate high-quality and relevant queries to represent each document. Query generation has been the topic of a number of recent works in the field of Information Retrieval.

A common example is docT5query (Nogueira and Lin, 2019), a neural document expansion technique that generates relevant queries and appends them to each document in the collection. Then, BM25 is used to perform retrieval on the augmented collection. This simple method can significantly improve on BM25. A follow-up study shows that even completely discarding the original document text and only using the generated queries to represent the documents can achieve better retrieval effectiveness than using the original document text (Lin et al., 2021b).

The TILDEv2 model, an effective and efficient sparse neural retrieval model, also uses document expansion based on query generation (Zhuang and Zuccon, 2021b). While one of the query generation methods adopted in TILDEv2 is docT5query, Zhuang&Zuccon have shown how the TILDE (Zhuang and Zuccon, 2021c) retrieval model can be exploited as a lightweight query generator. The use of TILDE in place of docT5query leads to similar retrieval effectiveness than docT5query but it requires several order of magnitude less computations (Zhuang and Zuccon, 2021b; Scells et al., 2022). The query generation method we use in DSI-QG is akin to docT5query. While the use of TILDE in place of docT5query for the query generation step of DSI-QG may be attractive because of its lower computational costs, we note that TILDE produces query terms that are independent of each other and thus is unlikely to be effective in for DSI-QG. In other words: TILDE generates unigram query tokens, not queries (i.e. sequences of tokens) – and these then are not representative of the inputs that the model will observe at retrieval time (e.g., Table 4 shows that queries in the two datasets considered in our work consists of, on average, 12.1 and 14.8 query tokens, respectively).

Query generation has also been used for the task of domain adaption and for generating training data for the zero-shot setting. Wang et al. proposed GPL, a framework for training domain adaptation rankers by generating pseudo labels with a query generation model  (Wang et al., 2021b). Bonifacio et al. directly used the GPT-3 model (Brown et al., 2020) to generate queries for training rankers in the few-shot setting (Bonifacio et al., 2022). Luo et al. introduced a domain-relevant template-based query generation approach which uses a sequence-to-sequence model conditioned on the templates to generate a large number of domain-related queries in a bid to mitigate the train-test overlap issue (Luo et al., 2022).

These prior works only focus on the mono-lingual retrieval setting. In contrast, our work also explores the usability of the query generation model for the cross-lingual information retrieval task.