RocketQAv2: A Joint Training Method for Dense Passage Retrieval and Passage Re-ranking
This August 2021 paper presents a process of optimising dense passage retrieval and passage re-ranking in natural language processing tasks.
Passage retrieval and re-ranking are essential components in various information retrieval tasks, such as question answering and document search.
The retriever aims to efficiently identify a subset of relevant passages from a large corpus, while the re-ranker refines the retrieved results to improve the relevance of top-ranked passages.
The paper recognises that both procedures contribute to the final performance, and thus, it is essential to optimise them together to achieve mutual improvement.
The approach includes two key components
Dynamic listwise distillation
The authors propose a unified listwise training approach for both the retriever and the re-ranker.
Listwise training is an approach where the model considers the entire list of candidate passages as a whole, rather than treating each passage independently.
In listwise training, the model learns to optimise the ranking of the entire list of passages based on their relevance scores.
During the dynamic distillation process, the retriever and the re-ranker can adaptively learn from each other's relevance information, which is crucial for mutual improvement in joint training. This is a significant advancement over previous distillation methods that typically froze one module while training the other.
Hybrid data augmentation
To support the listwise training approach, the authors design a hybrid data augmentation strategy that generates diverse training instances.
This is important because the listwise training approach requires diverse and high-quality training instances that can better represent the distribution of all the passages in the entire collection.
The paper's approach is interesting for several reasons:
Joint optimisation: By jointly training the retriever and the re-ranker, the approach enables the two modules to learn from each other and improve together.
Listwise training: The unified listwise training approach for both modules allows for a more consistent and effective learning process. It also enables the transfer of relevance information between the two modules during the dynamic distillation process.
Data augmentation: The hybrid data augmentation strategy addresses the need for diverse and representative training instances in the listwise training approach. This is crucial for the model to learn effectively and generalize well to unseen data.
Experimental results: The authors conduct extensive experiments on two benchmark datasets, MSMARCO and Natural Questions, demonstrating the effectiveness of their approach.
Note: Dense Passage Retrieval vs. Sparse Retrieval
Dense passage retrieval has been shown to outperform older sparse retrieval techniques.
Sparse methods rely on term frequency and document frequency to rank documents, which can miss nuanced contextual relationships between words.
Dense retrieval methods, by contrast, use deep learning to understand and utilise these subtleties, allowing them to retrieve passages that are more contextually relevant to the query.
Detailed Analysis of Methodology
The methodology of RocketQAv2 focuses on jointly training the dense passage retriever and the re-ranker to achieve mutual improvement.
The key components of the methodology are as follows:
Dynamic Listwise Distillation
RocketQAv2 introduces a unified listwise training approach for both the retriever and the re-ranker.
Given a query and a list of candidate passages, the retriever and re-ranker generate relevance scores for each passage.
The relevance scores are normalised in a listwise manner to obtain relevance distributions over the candidate passages.
What does listwise manner mean?
In the context of RocketQAv2, the relevance scores generated by the retriever and re-ranker are normalised in a listwise manner. This means that the relevance scores are transformed into probability distributions over the list of candidate passages, considering the scores of all passages in the list.
The KL-divergence between the relevance distributions from the retriever and the re-ranker is minimised to adaptively reduce the difference between them.
What is KL-divergence?
KL-divergence (Kullback-Leibler divergence) is a measure of the difference between two probability distributions. In RocketQAv2, the KL-divergence is used to quantify the difference between the relevance distributions generated by the retriever and the re-ranker. By minimising the KL-divergence, the model aims to make the relevance distributions from the retriever and re-ranker more similar, promoting consistency between them.
The re-ranker is also optimised using a cross-entropy loss based on the ground-truth labels of the passages.
Optimising with cross-entropy loss
Cross-entropy loss is a commonly used loss function for classification or ranking tasks. In RocketQAv2, the re-ranker is optimised using a cross-entropy loss based on the ground-truth labels of the passages.
The cross-entropy loss encourages the re-ranker to assign higher relevance scores to the positive (relevant) passages and lower scores to the negative (irrelevant) passages. The final loss function combines the KL-divergence loss and the supervised cross-entropy loss.
During distillation, both the retriever and the re-ranker are dynamically updated to adaptively synchronise and improve each other.
Hybrid Data Augmentation
To generate diverse and high-quality training instances for listwise training, RocketQAv2 employs a hybrid data augmentation strategy.
The RocketQA retriever is used to retrieve top-n passages from the corpus.
For undenoised instances, hard negatives are randomly sampled from the retrieved passages, and ground-truth positives are included.
For denoised instances, the RocketQA re-ranker is used to remove predicted negatives with low confidence scores and include denoised positives with high confidence scores.
The hybrid approach incorporates both undenoised and denoised instances to improve the diversity of the candidate passage list.
Undenoised and denoised instances
Undenoised instances refer to the training instances that are directly obtained from the retrieval results without any further processing.
Denoised instances, on the other hand, are obtained by applying a denoising step using the RocketQA re-ranker.
In the denoising step, the re-ranker is used to remove predicted negative passages with low confidence scores and include predicted positive passages with high confidence scores.
The inclusion of both undenoised and denoised instances in the training data helps improve the diversity and quality of the candidate passage list.
Training Procedure
The retriever and re-ranker are initialised with the learned dual-encoder and cross-encoder of RocketQA.
Initializing retriever and re-ranker with learned dual-encoder and cross-encoder
In RocketQAv2, the retriever and re-ranker are initialised using the pre-trained models from RocketQA.
The retriever is initialised with the learned dual-encoder model of RocketQA, which consists of two separate encoders for queries and passages.
The re-ranker is initialised with the learned cross-encoder model of RocketQA, which takes the concatenation of the query and passage as input.
By initialising the retriever and re-ranker with pre-trained models, RocketQAv2 leverages the knowledge and representations learned from RocketQA as a starting point for further training and optimisation.
The RocketQA retriever and re-ranker are used to generate training data via hybrid data augmentation.
Dynamic listwise distillation is performed to jointly optimise the retriever and re-ranker based on the final retrieval performance.
After training, the retriever and re-ranker can be applied for inference in a pipeline.
Key points and advantages of RocketQAv2
Joint training of the retriever and re-ranker using a unified listwise approach enables mutual improvement and adaptation.
Dynamic listwise distillation utilises soft relevance labels instead of hard pseudo-labels, allowing for more effective knowledge transfer between the modules.
Hybrid data augmentation incorporates both undenoised and denoised instances, improving the diversity and quality of training data.
The use of the KL-divergence loss and the supervised cross-entropy loss provides a balanced optimization objective for the joint training.
The methodology simplifies the training process and opens up the possibility for end-to-end training of the entire dense retrieval architecture.
RocketQAv2 builds upon the architecture and training techniques of RocketQA but introduces significant improvements through joint training, dynamic listwise distillation, and hybrid data augmentation. These advancements aim to enhance the retrieval performance and enable more effective collaboration between the retriever and the re-ranker.
Results
Passage Retrieval Results
RocketQAv2 and PAIR are two methods that perform significantly better than other methods in finding relevant passages from a large collection of text data.
RocketQAv2 is particularly good at ranking the most relevant passages at the top of the search results, likely due to its ability to learn from and mimic the behavior of a more advanced component called the re-ranker.
The choice of the underlying language model, which is a pre-trained component used to understand and process text, is not the only factor contributing to RocketQAv2's superior performance. Even when using the same language model as another method (DPR-E), RocketQAv2 still performs significantly better.
Passage Re-ranking Results
RocketQAv2's re-ranker, which is a component that refines the ranking of retrieved passages, significantly outperforms all other methods tested on the MSMARCO dataset.
The re-ranker's superior performance can be attributed to the joint learning process, where it is trained to closely match the relevance distribution of the retriever component using a technique called dynamic listwise distillation.
When the RocketQA re-ranker is replaced with the RocketQAv2 re-ranker, there is a notable improvement in the retrieval results, highlighting the effectiveness of the RocketQAv2 re-ranker.
Conclusion
RocketQAv2 is an approach to enhance the performance of passage retrieval and re-ranking systems through joint training and dynamic distillation.
By introducing a unified listwise training approach and a hybrid data augmentation strategy, RocketQAv2 enables the retriever and re-ranker to learn from each other and adapt their performance dynamically.
Extensive experiments on benchmark datasets demonstrate the effectiveness of the approach, outperforming state-of-the-art methods in both passage retrieval and re-ranking tasks.
The proposed techniques in RocketQAv2 have the potential to benefit a wide range of information retrieval applications, such as question answering, document search, and recommendation systems.
Future work may explore the application of our approach to other domains and investigate the integration of additional techniques, such as adversarial training and multi-task learning, to further improve the performance of retrieval and re-ranking systems.
Last updated