HongChien Yu, Chenyan Xiong, Jamie Callan
This repository holds code that reproduces results reported in Improving Query Representations for Dense Retrieval with Pseudo Relevance Feedback.
Dense retrieval systems conduct first-stage retrieval using embedded representations and simple similarity metrics to match a query to documents. Its effectiveness depends on encoded embeddings to capture the semantics of queries and documents, a challenging task due to the shortness and ambiguity of search queries. This paper proposes ANCE-PRF, a new query encoder that uses pseudo relevance feedback (PRF) to improve query representations for dense retrieval. ANCE-PRF uses a BERT encoder that consumes the query and the top retrieved documents from a dense retrieval model, ANCE, and it learns to produce better query embeddings directly from relevance labels. It also keeps the document index unchanged to reduce overhead. ANCE-PRF significantly outperforms ANCE and other recent dense retrieval systems on several datasets. Analysis shows that the PRF encoder effectively captures the relevant and complementary information from PRF documents, while ignoring the noise with its learned attention mechanism.
pip install -r requirements.txt
Run the following script to preprocess data:
cd data_prep
bash download_data.sh
bash preprocess_data.sh
cd data_prep
bash get_ance_embs.sh
bash get_ance_ranking.sh
Run the following command to create PRF data from ANCE top-retrieved documents:
cd data_prep
bash get_prf_data.sh
bash train_encoder.sh
While training is running, concurrently run
bash eval.sh
which keeps looking for the newest checkpoints and evaluate it on marco. This is sadly not a very effective use of GPU in terms of utilization percentage, but it makes the training faster by avoiding periodic switching from training to evaluation.
In our work, we picked the model that performs best on marco dev as reported by eval.sh tensorboard.
Trained models for k=3 can be downloaded here.
Ranking files for k=3 can be downloaded here.