Bridging the Training-Inference Gap for Dense Phrase Retrieval
Abstract
Building dense retrievers requires a series of standard procedures, including training and validating neural models and creating indexes for efficient search.
However, these procedures are often misaligned in that training objectives do not exactly reflect the retrieval scenario at inference time.
In this paper, we explore how the gap between training and inference in dense retrieval can be reduced, focusing on dense phrase retrieval (Lee et al., 2021a) where billions of representations are indexed at inference.
Since validating every dense retriever with a large-scale index is practically infeasible, we propose an efficient way of validating dense retrievers using a small subset of the entire corpus.
This allows us to validate various training strategies including unifying contrastive loss terms and using hard negatives for phrase retrieval, which largely reduces the training-inference discrepancy.
As a result, we improve phrase retrieval by 2-3% in top-1 accuracy and passage retrieval by 2-4% in top-20 accuracy for open-domain question answering.
Our work urges modeling dense retrievers with careful consideration of training and inference via efficient validation while advancing phrase retrieval as a general solution for dense retrieval.
1. Introduction
Dense retrieval aims to learn effective representations of queries and documents by making representations of relevant query-document pairs to be similar (Chopra et al., 2005; Van den Oord et al., 2018).
With the success of dense passage retrieval for open-domain question answering (QA) (Lee et al., 2019; Karpukhin et al., 2020), recent studies build an index for a finer granularity such as dense phrase retrieval (Lee et al., 2021a), which largely improves the computational efficiency of open-domain QA by replacing the retriever-reader model (Chen et al., 2017) with a retriever-only model (Seo et al., 2019; Lewis et al., 2021).
Building a dense retrieval system involves multiple steps (Figure 1) including training a dual encoder (§3), selecting the best model with validation (§2), and constructing an index (often with filtering (§A)) for an efficient search.
However, these components are somewhat loosely connected to each other.
For example, model training is not directly optimizing the retrieval performance on the full corpus on which models should be evaluated.
In this paper, we aim to minimize the gap between training and inference of dense retrievers to achieve better retrieval performance.
However, developing a better dense retriever requires validation, which requires building large indexes from a full corpus (e.g., the entire Wikipedia for open-domain QA) for inference with a huge amount of computational resources and time.
To tackle this problem, we first propose an efficient way of validating dense retrievers without building large-scale indexes.
Analysis of using a smaller random corpus with different sizes for the validation reveals that the accuracy from small indexes does not necessarily correlate well with the retrieval accuracy on the full index.
As an alternative, we construct a compact corpus using a pre-trained dense retriever so that validation on this corpus better correlates well with the retrieval on the full scale while keeping the size of the corpus as small as possible to perform efficient validation.
With our efficient validation, we revisit the training method of dense phrase retrieval (Lee et al., 2021a;b), a general framework for retrieving different granularities of texts such as phrases, passages, and documents.
We reduce the training-inference discrepancy by unifying previous loss terms to discriminate a gold answer phrase from other negative phrases altogether instead of applying in-passage negatives (Lee et al., 2021b) and in-batch negatives separately.
To better approximate the retrieval at inference where the number of negatives is extremely large, we use all available negative phrases from training passages and put more weights on negative phrases.
We also leverage model-based hard negatives (Xiong et al., 2020) for phrase retrieval, which hasn’t been explored in previous studies. This enables our dense retrievers correct mistakes made at inference time.
In summary, we introduce an efficient way of validation for dense retrievers to confirm and accelerate better modeling of dense retrievers and improve dense phrase retrieval [Figure] models with modified training objectives and hard negatives based on our efficient validation.
Consequently, we achieve 44.0% exact match score and 83.8% top-20 passage retrieval accuracy on Natural Questions (Kwiatkowski et al., 2019) and 55.5% exact match score and 82.7%
top-20 passage retrieval accuracy on TriviaQA (Joshi et al., 2017), the state-of-the-art retrieval-only open-domain QA accuracy.
2. Efficient Validation of Phrase Retrieval
Our goal is to train a dense retriever M that can accurately find a correct answer in the entire corpus C (in our case, Wikipedia).
Careful validation is necessary to confirm whether new training methods are truly effective. It also helps finding optimal configurations induced by those techniques.
However, building a large-scale index for every model makes the model development process slow and also requires huge memory.
Thus, an efficient validation method could expedite modeling innovations in the correct directions.
It could also allow frequent comparison of different checkpoints when updating a full index simultaneously during the training is computationally infeasible.
Measuring the retrieval accuracy on an index from a smaller subset of the full corpus (denoted as C ⋆ ) for model validation would be a practical choice, hoping argmax M∈ Ω acc
( D |M , C ⋆ ) ≈ argmax M∈ Ω acc ( D |M , C ) where Ω is a set of model candidates and acc means the retrieval accuracy on a QA dataset D . We first examine how a relative order of accuracy between modeling approaches may change with varying sizes of the random subcorpus (§2.1) and then develop a clever way to construct a compact subcorpus that maintains reasonable correlation with the retrieval accuracy in the full scale (§2.2). 2.1. Random Subcorpus Reading comprehension (RC) can be regarded a special case of open-domain QA, where a corpus contains only a single gold passage (i.e., C q = { c } ) for each question.
Here, the subcorpus is question-dependent.
We first gather all gold passages from the development set as a small corpus C 0 , a minimal set that contains answers to all development set questions.
We consider a corpus R r whose size is r times the size of the full corpus by simply appending C 0 with random passages by sampling from the full corpus C , i.e., C {+
+} ⊂ R r ⊂ C and | R r | = r | C | .
We first gather all gold passages from the development set as a small corpus C 0 , a minimal set that contains answers to all development set questions.
We consider a corpus C r whose size is r times
− log e s ( q,p + ; c ) e s ( q,p + ; c ) + (cid:80) ( p - ; c ) ∈ N inp e s ( q,p - ; c ) − λ log e s ( q,p + ; c ) e s ( q,p + ; c ) + (cid:80) ( p - ; c ′ ) ∈ N inb ∪ N prb e s ( q,p - ; c ′ )
(1) − log e s ( q,p + ; c ) e s ( q,p + ; c ) + (cid:80) ( p - ; c ′ ) ∈ N inp ∪ N inb ∪ N prb ∪ N hard λ ( p - ) e s ( q,p - ; c ′ )
(2) the size of the full corpus by simply appending C 0 with random passages by sampling from the full corpus C , i.e.,
C 0 ⊂ C r ⊂ C and | C r | = r | C | .
We specifically use r = 1 / 100 , 1 / 10 in our experiments.
As the corpus size increases, finding the correct answer from a larger number of possible candidates becomes more difficult, so the retrieval accuracy generally decreases (Reimers & Gurevych, 2021).
DensePhrases (Lee et al., 2021a) simply choose the best checkpoint with the highest RC accuracy assuming that a model with better RC accuracy leads to a better retrieval accuracy, or use the last checkpoint at the end of the training.
It is problematic since our preliminary experiments demonstrate that the RC accuracy and the retrieval accuracy on different sizes of corpus including the full corpus, do not necessarily correlate well with each other.
Using a large subcorpus is better for accurate validation not to deviate much from the trends of retrieval accuracy of a full corpus.
However, a smaller subcorpus would be better in terms of validation efficiency. This trade-off drives us to design a better way of constructing a validation corpus. 2.2. Hard Subcorpus The retrieval accuracy given a subcorpus C ⋆ should have a high correlation with the retrieval accuracy over the full corpus and the size of corpus | C ⋆ | should be small enough (or as small as possible) for efficient validation.
For a reasonably accurate dense retriever, it is relatively easy to discriminate a gold phrase from other phrases in random passages.
Therefore, it is better to collect a subcorpus with hard passages to test dense retrievers on a similar condition to a full corpus which includes many difficult phrases to discriminate if the corpus can have a limited number of negative passages.
We construct a hard corpus H k with a compact size using a pre-trained retriever ¯ M to extract all context passages of top- k retrieved phrases for all query q in the development set Q dev , and C 0 is merged to always include an answer, i.e., H k = C 0 ∪ (cid:83) ( q,a ) ∈ Q dev ¯ M k ( q | C ) where M k ( q | C ) denotes the top- k passage retrieval results for a query q from the model M .
If ¯ M is reasonably accurate, negative examples retrieved by ¯ M will make our new model M difficult to find a correct answer. We expect the retrieval accuracy from
H k quickly drops as k increases and reaches close to the retrieval accuracy on the full corpus C with a manageable k so that we can use retrieval accuracy on a hard subcorpus for efficient validation.
3. Optimized Training of DensePhrases 3.1. Background: Training of DensePhrases
The question encoder and the phrase encoder are jointly trained using reading comprehension datasets.
A phrase p is represented as a concatenation of start and end token vectors from the contextualized representations of a context passage c using a phrase encoder.
A question q is represented as a concatenation of vectors using two different encoders for the start and the end.
The main training objective is a sum of the two separate contrastive loss terms weighted by the λ coefficient as formally defined in Equation 1.
2 One is for contrasting a phrase token of positive start/end position ( p + ) to that of other positions in the context passage ( N inp =
{ ( p ; c ) ̸
= ( p + ; c ) | p ∈ c } ).
Another is for contrasting the same token to other positive tokens in a current ( N inb =
{ ( p ′ ; c ′ )
̸
= ( p + ; c ) | ( q ′ , p ′ ;
c ′ ) ∈ B} ) or previous T minibatches ( N prb =
{ ( p ′ ; c ′ ) | ( q ′ , p ′ ; c ′ )
∈ B pre } ).
The numbers of negatives are | N inp | =
L − 1 , | N inb | =
B − 1 , and | N prb | =
B × T where L is the sequence length of context passages and B (= |B| ) is the batch size.
Additional training techniques for DensePhrases are described in Appendix B.
3.2. Unified Loss
The original training objective of DensePhrases (Equation 1) has separate terms for finding a relevant passage (in/prebatch negatives) and finding the exact phrase position in the passage (in-passage negatives).
However, we should find an answer phrase among all possible candidates at once during the test time.
Therefore, we modify the loss term as a unified version (Equation 2) by putting all negatives together into the contrastive targets.
We also introduce the λ coefficient to the unified loss to give weights to negatives.
It opens a new question of how we should set the value of λ .
The role of λ can be interpreted in two ways.
First, multiplying λ to an exponential of a score is equivalent to adding a positive value to the score ( λe s =
e s +log λ ), and then the loss term becomes the soft version of margin-based loss.
Second, using λ can mimic the inference time where the number of negative tokens is much larger by duplicating a negative λ times ( λe s =
e s + e s
+ ...
+ e s ) to close the gap between training and test.
Based on the second interpretation, we set different value of λ depending on where negative phrase p - is from: λ ( p - ) = λ inp δ ( p - ∈ N inp )
+ λ inb δ ( p - ∈ N inb ∪ N prb ) + λ hard δ ( p - ∈ N hard ) . We extend to use all tokens in context passages with a similar intuition that contrasting with as many tokens as possible could be helpful instead of using only start/end position tokens to in/pre-batch negatives.
It changes in/pre-batch negatives to N inb
=
{ ( p ; c ′ ) ̸
= ( p + ; c ) | p ∈ c ′ , ( q ′ , p ′ ; c ′ ) ∈
B} and N prb =
{ ( p ; c ′ )
| p ∈ c ′ , ( q ′ , p ′ ;
c ′ ) ∈ B pre } ) and their sizes | N inb | =
B × L − 1 and | N prb | =
B × T × L .
This change also increases the number of negatives hundreds of times and turns out empirically advantageous.
3.3. Hard Negatives for Phrase Retrieval
We exploit hard negatives to benefit phrase retrieval, a widely used technique for passage retrieval 3 but never fully examined for phrase retrieval.
We use a encoder model and phrase index from the first round to extract model-based hard negatives from top- k phrase retrieval results for questions in the training set.
We exclude examples when a context passage of a retrieved phrase contains an answer.
A context passage corresponding to a retrieved phrase can be restored using information stored along with the index.
It helps to focus more on topically different documents and shares the intuition from the analysis in Lee et al.
(2021b) that DensePhrases less rely on the topical relevance than DPR.
Using high-quality hard negatives by removing false negatives is important to train a better model.
We left filtering based on a cross encoder (Qu et al., 2021;
Ren et al., 2021) to future work due to the convenience of automatic filtering.
After the hard negative mining, we fine-tune a dual encoder with the hard negatives.
We sample h hard negatives for each training step and append them to negative targets for the loss calculation.
We expect that hard negatives give a better training signal than random in/pre-batch negatives (Xiong et al., 2020) because those are examples difficult to discriminate for the previous model.
Moreover, hard negatives extracted from the larger corpus could expose a model to other diverse documents than the original training dataset.
This is similar to query-side fine-tuning but differs in that both the question encoder and phrase encoder are updated.
Similar to §3.2, we include all tokens in the context passage of all hard negatives in a mini-batch for N hard .
Using all available tokens is generally better because they potentially belong to the final negative phrase candidates in inference.
Training with larger numbers of negatives is beneficial to reduce the gap between training and inference.
Including all of them does not induce significant additional memory overhead since we should encode the same number ( h ) of passages regardless of different options.
Therefore, we use N hard =
(cid:83) (ˆ p ;
ˆ c ) ∈ H ( q,p + ; c ) , ( q,p + ; c ) ∈B { ( p ; ˆ c )
| p ∈
ˆ c } as all tokens from all hard negatives in a mini-batch where H is a set of the sampled h hard negatives and | N hard | =
B × L × h .
3.4. Token Filter
We modify to train a token filter which is used to reduce the index size separately rather than joint training with the dual encoder.
Appendix A contains a more detailed explanation and interesting analysis of the filter threshold.
4. Experiments
To show the effectiveness of our proposed method, we evaluate DensePhrases models on open-domain QA benchmarks following the experimental setup of Lee et al.
(2021a;b). Datasets We measure phrase retrieval accuracy and passage retrieval accuracy on two open-domain QA datasets: Natural Questions (Kwiatkowski et al., 2019) and TriviaQA (Joshi et al., 2017).
We train our phrase retrieval models with two different datasets followed by query-side fine-tuning on a target dataset: DensePhrases ♡ with Natural Questions reading comprehension dataset and DensePhrases ♠ with a combination of multiple reading comprehension datasets.
Training details are described in Appendix C. We build the phrase index using the 2018-12-20 Wikipedia snapshot ( C ) and its different smaller subsets C ⋆ .
Model Validation
In our preliminary experiments, we observe that the best checkpoint among training epochs differs depending on the corpus size (especially for small scale). Figure 2 shows validation retrieval accuracy of the
DensePhrases models with different training methods on various sizes of random and hard subcorpora. The retrieval accuracy on the hard subcorpus rapidly drops and reaches close to the retrieval accuracy on the full corpus as k increases with moderately increasing the index size. On the other hand, retrieval accuracy on a random subcorpus is higher than on a hard subcorpus with similar index size. For instance, retrieval accuracies on H 8 (5.1M) are lower than
[Table]
♡ -UL-HN 44. DensePhrases
[Figure]
-UL 42.
[Figure] [Figure]
[Figure]
e t r i e v a l A cc u r a c y ( E M @ 1 )
[Figure]
those on R 1 / 100 (24.2M) with 4 times smaller index, and retrieval accuracies on H 16 (8.7M) are lower than those on
R 1 / 10 (266.4M) with 30 times smaller index.
It indicates that a hard subcorpus can effectively imitate inference with a full corpus where correct retrieval is the most difficult.
The relative order of accuracy between models on hard subcorpus converges quickly at around H 10 (6.1M).
Contrarily, the order changes from R 1 / 100 to R 1 / 10 , meaning that it is difficult to be certain about the superiority of a model over another one based on a random subcorpus with an even larger index.
Therefore, we can use retrieval accuracy on a hard subcorpus as an efficient validation metric.
Validation results on any subcorpus clearly demonstrate that unified loss is helpful.
Query-side fine-tuning decreases
RC accuracy and retrieval accuracy in C 0 (1.1M) while improving retrieval accuracy in larger indexes.
It proves that a hasty conclusion with a small corpus size may avoid modeling development.
Open-domain QA results Table 1 summarizes experimental results on open-domain QA phrase retrieval and passage retrieval.
Both unified loss and hard negatives (Appendix D) are shown to be effective.
With the better training methods, our best model improves exact match accuracy on phrase retrieval from the original DensePhrases model by 2.7% in Natural Questions and 2.0% in TriviaQA, achieving the state-of-the-art retrieval-only performance.
Our methods also improve passage retrieval accuracy significantly.
Our best model improves top-20 accuracy on passage retrieval from the original DensePhrases model by 4.0% in Natural Questions and 1.8% in TriviaQA.
We may use DensePhrases as a building block of other tasks and expect to achieve good phrase retrieval performance with expressive reader models like FiD (Izacard & Grave, 2021).
5. Conclusion
In this study, we aim to bridge the gap between training and inference of phrase retrieval.
We first develop an efficient validation metric that measures retrieval accuracy on the index from a small corpus with hard passages using a pre-trained retriever.
Based on this validation, we show that the improvements in training of dense phrase retrieval with unified loss and hard negatives are effective.
As a result, we achieve state-of-the-art phrase retrieval and passage retrieval accuracy in open-domain question answering among retrieval-only approaches.
We describe related work in Appendix E and potential implications in Appendix F.
A. Token Filtering
Representation filtering is often applied in practice to reduce the index size for efficient search (Min et al., 2021).
For phrase retrieval, tokens that are not likely to be a start/end position of an answer are filtered out using a trained filter classifier based on a logit score for each token.
Only tokens with a score larger than a specific threshold are kept.
After the filtering, the index is compressed using optimized product quantization (Ge et al., 2013).
A.1. Token Filter Threshold
A filter threshold for the token filter determines a trade-off between the index size (efficiency) and retrieval accuracy (Figure 3).
Interestingly, we find that token filtering can even improve retrieval accuracy.
As we increase a threshold from a very small value (not filtering), the accuracy fluctuates but generally increases until a specific threshold because the filter successfully reduces the number of candidates, making prediction easier.
After that threshold, the accuracy drops quickly because the filter starts to leave out many correct tokens.
However, finding the peak retrieval accuracy requires a manual search of different thresholds after indexing and evaluating.
Since using it as a validation metric is expensive, we first select the best checkpoint based on retrieval accuracy without performing any token filtering.
Especially when the token filter is in the middle of training, the peak threshold will vary, and using a specific fixed threshold would not be fair.
Also, the best threshold changes depending on the corpus size, so choosing a threshold for the full corpus based on a smaller corpus is not straightforward.
A.2. Token Filter Training and Valiation
Lee et al.
(2021a) jointly train a token filter classifier with a dual encoder.
It is convenient in that an additional training process is not required, while we should tune on the weight for a loss before adding to the overall training loss.
Training pushes phrase vectors into two moving corns toward the start and end vectors since a logit is a dot product score between a phrase vector and a start/end vector.
It has two potential disadvantages: (1) phrase representations are concentrated on the part of the entire feature space, so the expressiveness of the model is not fully exploited, and (2) optimization is more difficult because of the moving targets.
To address the issues, we change to train a token filter after training a phrase encoder.
We could expect that the twostage training process encourages phrase representations to be distributed over the space.
Moreover, we can validate the token filter separately due to the separate training process and pick the best one.
We can not decide the threshold during the filter training, so we use the AUC-PR metric for filter validation by measuring precision and recall by sweeping all threshold values.
B. Training Techniques for DensePhrases
To learn better representations, the dual encoder is first pretrained with question-answer pairs generated by a question generation model as a data augmentation mainly for better reading comprehension capability and then fine-tuned with original question-answer pairs.
Also, knowledge distillation (Hinton et al., 2015) from a stronger reading comprehension model based on a cross encoder to the dual encoder is performed.
Documents in the reading comprehension dataset used for the training take only a tiny portion of the entire Wikipedia, and only a small number of negatives for each questionphrase pair are contrasted compared to billion-scale possible phrase candidates in the test time.
The query encoder can be further fine-tuned to reduce this discrepancy between training and inference while fixing the phrase encoder and the index by maximizing the likelihood of the gold answer among retrieved phrases for each question.
This process is called query-side fine-tuning .
C. Training details
We use almost the same training hyperparameters of the original DensePhrases except for the batch size B = 48 .
We use the number of training epochs to 2 with the generated question-answer pairs and increase the number of training epochs to 10 with the standard reading comprehension dataset for more careful validation.
We set λ ( p - ) to 256 for p - ∈ N inb ∪ N prb , otherwise to 1.
We set k =and h = 1 for the hard negative mining and sampling.
Our token filter achieves a much higher AUC-PR value than the filter of the original DensePhrases model (e.g., 0.348 vs.
We use a filter threshold of -3 for the index with the full corpus.
D. Discussion on Hard Negatives
With hard subcorpora, a model trained with hard negatives shows higher validation accuracy than a model without hard negative training before query-side fine-tuning, but the order changes after query-side fine-tuning (Figure 2).
We blame it on the similarity between the hard corpus construction and the hard negative mining process.
Therefore, we may have to rely on another metric (including retrieval accuracy on random subcorpus) only for validating the effectiveness of hard negatives.
Interestingly, hard negatives improve in-domain accuracy but decrease out-of-domain accuracy.
Since gathered hard negative passages are close to the original training data, it
[Figure]
R e t r i e v a l A cc u r a c y ( E M @ 1 )
[Figure]
e t r i e v
[Figure] [Figure] [Figure] improve the performance of questions from the same domain but could overfit and may lose the generalization ability.
This observation solicits better hard negative mining methods.
E. Related Work
Dense retrieval Retrieving relevant documents for a query (Mitra & Craswell, 2018) is crucial in many NLP applications like open-domain question answering and knowledge-intensive tasks (Petroni et al., 2021).
Dense retrievers typically build a search index for all documents by pre-computing the dense representations of documents using an encoder.
Off-the-shelf libraries for a maximum inner product search (MIPS)
(Johnson et al., 2019; Guo et al., 2020) enable model training and indexing to be developed independently (Lin, 2022).
However, both training dense retrievers and building indexes should take into account the final retrieval accuracy.
In this respect, we aim to close the gap between training and inference of dense retrievers.
Phrase retrieval Phrase retrieval (Seo et al., 2019) directly finds an answer with MIPS from an index of contextualized phrase vectors.
This removes the need to run an expensive reader for open-domain QA.
As a result, phrase retrieval allows real-time search tens of times faster than retriever-reader approaches as an alternative for opendomain QA.
DensePhrases (Lee et al., 2021a) removes the requirement of sparse features and significantly improves the accuracy from previous phrase retrieval methods (Seo et al., 2019; Lee et al., 2020).
Lee et al.
(2021b) show how retrieving phrases could be translated into retrieving larger units of texts like a sentence, passage, or document, making phrase retrieval a general framework for retrieval.
Despite these advantages, phrase retrieval requires building a large index from billions of representations.
In this work, we focus on improving phrase retrieval with more efficient validation.
Validation of dense retrieval Careful validation is essential for developing machine learning models to find a better configuration (Melis et al., 2018) or avoid falling to a wrong conclusion.
However, many works on dense retrieval do not clearly state the validation strategy, and most of them presumably perform validation on the entire corpus.
It is doable but quite expensive 7 to perform frequent validation and comprehensive tuning.
Hence, it motivates us to devise efficient validation for dense retrieval.
Like ours, Hofst ¨ atter et al.
(2021) construct a small validation set by sampling queries and using a baseline model for approximate dense passage retrieval but limited to early stopping.
Hard examples Adversarial data collection by an iterative model (or human) in the loop process aims to evaluate or reinforce models’ weaknesses, including the robustness to adversarial attacks (Kaushik et al., 2021;
Bartolo et al., 2021; Nie et al., 2020; Kiela et al., 2021).
In this work, we construct a compact corpus from a pre-trained dense retriever for efficient validation.
Also, we extract hard negatives from retrieval results of the previous model for better dense representations.
F. Broader Impact
Our work demonstrates that thorough validation is crucial for the accurate and efficient development of phrase retrieval with large corpus.
Also, it could be especially beneficial in real applications where the corpus size is much larger than benchmark datasets.
Moreover, we prove that modeling and training methods should be designed closely to retrieval in inference time
Despite its remarkable efficiency and flexi- bility, phrase retrieval has been relatively less studied than passage retrieval.
We believe that our work can encourage more study on phrase retrieval with an efficient development cycle.
Furthermore, we hope that our findings could be extended to dense retrieval in general to help a wide variety of applications.