Advanced Information Retrieval - Final Project

Query Adaptive Contextual Document Embeddings

Technion - Israel Institute of Technology
Faculty of Data Science and Decisions

---

Table of Contents ⚙️

1. About the Article
2. Project Overview
3. About the code
4. Running Instructions
5. Results & Comments

---

About the article

Our project is an extension on the following article:

Contextual Document Embeddings by John X. Morris and Alexander M. Rush.

The article "Contextual Document Embeddings" by John X. Morris and Alexander M. Rush from Cornell University proposes a method to improve document embeddings by incorporating context from surrounding documents. Traditional embeddings, which rely solely on the individual document, can be insufficient for highly specific information retrieval tasks. This paper addresses this by incorporating neighboring documents into the embedding process, making the representation more context-aware. The authors introduce a contextual training objective and an architecture that explicitly encodes information from neighboring documents, demonstrating improved performance in various scenarios.

Project Overview

In this project, we aimed to extend the "Contextual Document Embeddings" model by introducing a query-adaptive approach. While the original model effectively contextualizes document embeddings by considering neighboring documents, it does not account for the user's needs expressed in the query, which is essential for information retrieval tasks.

To achieve this, we developed two models. The first model, called Query Adaptive Contextual Document Embedder (QACDE), makes the document embedding query-adaptive, meaning that the document embedding is influenced by the query's embedding, following the same embedding method proposed in the original paper. However, upon reflection, we realized that the document embedding greatly depends on how the query is represented. This led us to develop a second model called Multi-Embeddings-Query Adaptive Contextual Document Embedder (MEQCDE), that computes the document embedding using multiple query embeddings, such as BERT, TF-IDF and query embedding generated by the model of the article. The goal of this approach is to combine the strengths of different query embedding methods to create more robust and informative document embeddings. By leveraging multiple representations of the query, we aim to capture various aspects of the document-query relationship, improving the relevance of the information retrieval process.

About the code

The code for this project is built upon the original implementation from the paper. We first use the pre-trained model from the article to generate document and query embeddings. These embeddings are then utilized to train our own models. Our code handles the training and evaluation of both our query-adaptive models as well as the evaluation of the original model from the paper (which serves as our baseline). Additionally, it manages the query embedding generation using TF-IDF and BERT methods. The evaluation is conducted using Mean Average Precision (MAP), with the goal of comparing the performance of our models and the baseline across various BeIR datasets such as SciFact, FiQA, and NFCorpus. The training procedure includes contrastive negative sampling and a custom Loss function tailored to the specific training conditions, ensuring a thorough and meaningful comparison between the models.

Running Instructions

This repository contains the code and instructions to run the code. It is preferable to have acces to a Virtual Machine with a GPU because of the computational cost of the models.

1. If you have not VM, skip this step and go directly to step 2. Connect to the VM. Start the VM and run the following code in the terminal:

ssh <user-name>@<host-name>

where is your username and is the public adress IP of the VM. Write your password to finish the connection.

2. Enter in the project file via the following command on the terminal:

cd query_adaptive_contextual_document_embedding

3. Install all the dependencies via the requirements.txt file. Run this command on the terminal:

pip install -r requirements.txt

4. Run the module Python run_pre_trained_cde.py to save into pickle files the document and queries embeddings by the model of the article (CDE):

python run_pre_trained_cde.py --dataset <dataset_name> --seed <seed>

where <dataset_name> is a correct name of a dataset present in the Benchmark BEIR and is the an integer that represents the seed for this module.

5. Run the module compute_query_embeddings.py to save into pickle files the query embeddings for different methods of embeddings.

python compute_query_embeddings.py --dataset <dataset_name> --tfidf_max_dim <max_dim_emb> --seed <seed>

where <dataset_name> is a correct name of a dataset present in the Benchmark BEIR, <max_dim_emb> represents the dimension maximal of the TF-IDF query embedding and is the integer that represents the seed of this module that has to be the equal to the seed of step 3.

6. Run the module baseline_run.py to obtain MAP results for the chosen dataset for the model of the article (CDE).

python baseline_run.py --dataset <dataset_name> --index <index_num>

where <dataset_name> is a correct name of a dataset present in the Benchmark BEIR and <index_num> is the seed we use so far in steps 4 & 5.

7. Run the run_query_adaptive_layer.py to obtain MAP results for the chosen dataset for the QACDE and MEQACDE models.

python run_query_adaptive_layer.py --dataset <dataset_name> --index <index_num> --hidden_dim <hidden_dim> --epochs <num_epochs> --seed <seed>

where <dataset_name> is a correct name of a dataset present in the Benchmark BEIR, <index_num> is the seed we used so far in step 4, 5 & 6, <hidden_dim> is the dimension of the hidden layer of the two models, <num_epochs> is the number of training epochs and is the seed of the program.

Results & Comments

For this part, we ran the two models with hidden_dim = 1024 and epochs = 20, for several seed to perform an averaged MAP for differents datasets.

The results clearly demonstrate the effectiveness of adding the query-adaptive layer. Both QACDE and MEQACDE consistently out-perform CDE baseline across all three datasets. The improvements, ranging from +3.0% to +4.7% relative MAP, show that dynamically adapting document embeddings based on the query provides benefits for retrieval accuracy. Furthermore, when comparing QACDE and MEQACDE models, an interesting observation arises. For smaller datasets like NFCorpus and SciFact, the simpler QACDE model slightly outperforms MEQACDE, suggesting that the extra query information in MEQACDE might not add significant value or may introduce complexity on smaller data scales. However, for the larger FiQA dataset, MEQACDE significantly outperforms QACDE, likely due to the larger data size allowing MEQACDE to better synthesize multiple query signals (CDE, BERT, TF-IDF). This suggests that MEQACDE is more effective on larger, diverse datasets, where different query embeddings provide complementary information.