Bridging the Training-Inference Gap for Dense Phrase Retrieval

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_{0}\subset R_{r}\subset C$ and $|R_{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 and 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.³³3RC accuracy generally improves during the training of DensePhrases as the training loss directly optimizes it. 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.

The retrieval accuracy given a subcorpus $C^{\star}$ should have a high correlation with the retrieval accuracy over the full corpus and the size of corpus $|C^{\star}|$ 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 $\bar{\mathcal{M}}$ to extract all context passages of top-$k$ retrieved phrases for all query $q$ in the development set $Q_{\text{dev}}$, and $C_{0}$ is merged to always include an answer, i.e., $H_{k}=C_{0}\cup\bigcup_{(q,a)\in Q_{\text{dev}}}\bar{\mathcal{M}}_{k}(q|C)$ where $\mathcal{M}_{k}(q|C)$ denotes the top-$k$ passage retrieval results for a query $q$ from the model $\mathcal{M}$. If $\bar{\mathcal{M}}$ is reasonably accurate, negative examples retrieved by $\bar{\mathcal{M}}$ will make our new model $\mathcal{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. It keeps the relative order of models with a much smaller size than the random subcorpus.

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 $\lambda$ coefficient as formally defined in Equation 1.⁴⁴4We denote the similarity score between a question $q$ and a phrase $p$ as $s(q,p;c)$. While the score and the loss term of start and end tokens are separately calculated in practice, we abbreviate it in the equation for simplicity. One is for contrasting a phrase token of positive start/end position ($p^{\text{+}}$) to that of other positions in the context passage ($N_{\text{inp}}=\{(p;c)\neq(p^{\text{+}};c)|p\in c\}$). Another is for contrasting the same token to other positive tokens in a current ($N_{\text{inb}}=\{(p^{\prime};c^{\prime})\neq(p^{\text{+}};c)|(q^{\prime},p^{\prime};c^{\prime})\in\mathcal{B}\}$) or previous $T$ mini-batches ($N_{\text{prb}}=\{(p^{\prime};c^{\prime})|(q^{\prime},p^{\prime};c^{\prime})\in\mathcal{B}_{\text{pre}}\}$). The numbers of negatives are $|N_{\text{inp}}|=L-1$, $|N_{\text{inb}}|=B-1$, and $|N_{\text{prb}}|=B\times T$ where $L$ is the sequence length of context passages and $B$ (= $|\mathcal{B}|$) is the batch size.

To learn better representations, the dual encoder is first pre-trained 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. Lastly, the token filtering classifier (explained more in §5) that discriminates tokens likely to be a start or end of the answer phrases using a linear classifier on top of phrase representations is jointly trained with the dual encoder. We omit two additional loss terms from Equation 1 for knowledge distillation and a token filtering classifier loss for brevity.

The original training objective of DensePhrases (Equation 1) has separate terms for finding a relevant passage (in/pre-batch 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 $\lambda$ coefficient to the unified loss to give weights to negatives. It opens a new question of how we should set the value of $\lambda$. The role of $\lambda$ can be interpreted in two ways. First, multiplying $\lambda$ to an exponential of a score is equivalent to adding a positive value to the score ($\lambda e^{s}=e^{s+\log\lambda}$), and then the loss term becomes the soft version of margin-based loss. Second, using $\lambda$ can mimic the inference time where the number of negative tokens is much larger by duplicating a negative $\lambda$ times ($\lambda 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 $\lambda$ depending on where negative phrase $p^{\text{-}}$ is from: $\lambda(p^{\text{-}})=\lambda_{\text{inp}}\delta(p^{\text{-}}\in N_{\text{inp}})+\lambda_{\text{inb}}\delta(p^{\text{-}}\in N_{\text{inb}}\cup N_{\text{prb}})+\lambda_{\text{hard}}\delta(p^{\text{-}}\in N_{\text{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_{\text{inb}}=\{(p;c^{\prime})\neq(p^{\text{+}};c)|p\in c^{\prime},(q^{\prime},p^{\prime};c^{\prime})\in\mathcal{B}\}$ and $N_{\text{prb}}=\{(p;c^{\prime})|p\in c^{\prime},(q^{\prime},p^{\prime};c^{\prime})\in\mathcal{B}_{\text{pre}}\}$) and their sizes $|N_{\text{inb}}|=B\times L-1$ and $|N_{\text{prb}}|=B\times T\times L$. This change also increases the number of negatives hundreds of times and turns out empirically advantageous.

We exploit hard negatives to benefit phrase retrieval, a widely used technique for passage retrieval⁵⁵5Karpukhin et al. (2020) use one hard negative obtained from BM25 per example in addition to in-batch negatives for training a dual encoder. Xiong et al. (2020) globally select hard negatives from the entire corpus with asynchronously index updates for faster convergence. RocketQA Qu et al. (2021) denoises hard negatives using cross encoder. The best strategy for hard negative mining and training is still an open problem in dense retrievals. but never fully examined for phrase retrieval. We perform a model-based hard negative mining by retrieving top phrases using a pre-trained dual encoder and an index built from this model. We filter out phrases whose surrounding passage includes a gold answer (§4.3.1) and then fine-tune the model with extracted hard negatives (§4.3.2). Although we do it only once, this process could be repeated until convergence.

A filter threshold for the token filter determines a trade-off between the index size (efficiency) and retrieval accuracy (Figure 2). 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.

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 cones 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 a subset 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 train a token filter after training a phrase encoder. We could expect that the two-stage 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.

In our preliminary experiments, we observe that the best checkpoint among training epochs differs depending on the corpus size (especially for small scale). Figure 9 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 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). However, the order of accuracy when using random subcorpus changes from $R_{1/100}$ to $R_{1/10}$ showing the difficulty of efficient validation. On the other hand, retrieval accuracy on a hard subcorpus is more stable and serves as an efficient validation metric.

Our validation results clearly demonstrate that unified loss is helpful. Query-side fine-tuning also harms the RC accuracy and the retrieval accuracy with $C_{0}$ (1.1M) while improving the retrieval accuracy with larger indexes. It shows how a wrong conclusion can be made from small-sized corpora.

Table 1 summarizes end-to-end open-domain QA results. Both unified loss and hard negatives are shown to be effective. With our improved training methods, the best model surpasses the original DensePhrases model by 2.7 points in Natural Questions and 2.0 points in TriviaQA, achieving the state-of-the-art retrieval-only open-domain QA performance.

Table 2 summarizes open-domain QA passage retrieval results. Our method also improves passage retrieval accuracy significantly. The best model improves top-20 passage retrieval accuracy by 4.0 points in Natural Questions and 1.8 points in TriviaQA. It again shows that DensePhrases can be used for passage retrieval as well. 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 and Grave, 2021).

From Footnote 9, we observe that with hard subcorpora, the model trained with hard negatives (cyan) shows higher validation accuracy than the model without hard negative training (yellow) before query-side fine-tuning, but their order changes after query-side fine-tuning (blue vs. red). This is because hard negative mining process is similar to hard corpus construction, blurring the accurate estimation of validation performance. However, we pick the best model before the query-side fine-tuning, which lets us to decide to go with hard negatives (due to cyan vs. yellow) and achieve state-of-the-art performance with the full index.

From Table 1, we observe that hard negatives improve in-domain accuracy but harm the out-of-domain accuracy. Since hard negative passages are close to the original training data, it improves the performance of questions from the same domain but could cause overfitting and harm the generalization ability. This observation solicits better hard negative mining methods.