Document Aware Contrastive Learning Approach for Generative Retrieval

Authors: maunendra  sankar desarkar
Venue: ACM_org_SIGIR_2024_Workshop_Gen-IR
Published: 2024-05-03 13:25:26

Abstract: The field of neural retrieval models has seen significant advancements with the rise of dense retrieval techniques, where documents are indexed based on embeddings from neural networks, particularly transformer-based models. Generative retrieval, a newer paradigm, employs sequence-to-sequence models to generate unique identifier strings for passages, providing an alternative approach to retrieval. However, these generative models primarily focus on learning the mappings between the queries and the identifiers, which may limit their ability to fully learn the relationship between queries and passages. This work introduces a novel method inspired by dense retrieval to enhance the learning of context-aware representations, thereby fostering a deeper understanding of the relationship between queries and documents. Additionally, a curriculum-based learning strategy is proposed to optimize contrastive losses effectively. Extensive experiments were conducted using the publicly available Natural Questions dataset to evaluate the proposed approach. The results demonstrate modest improvements in performance across all metrics, highlighting the robustness of the method.

Main paper: 
 introduction: 
 Information retrieval (IR) is a pivotal field focused on obtaining relevant information from extensive collections of data. By employing various methods and technologies, IR systems aim to locate and present information that meets users' needs from sources such as databases, web pages, text corpora, etc. Traditionally, IR methods have been employed in various applications, including web search engines, recommendation systems, open-domain question answering, and knowledge-grounded conversation [7] . In recent times, the importance of information retrieval has grown, as it forms the backbone of several emerging technologies such as retrieval-augmented generation [8, 13] , knowledge-intensive tasks [17] , and long-range transformers [3] .Sparse retrieval methods, such as those based on keyword matching and term frequency-inverse document frequency (TF-IDF), have long been foundational in information retrieval, offering quick search times and straightforward implementation. However, these methods struggle with capturing semantic relationships and context, treating terms independently and resulting in sub-optimal Permission to make digital or hard copies of all or part of this work for personal or classroom use is granted without fee provided that copies are not made or distributed for profit or commercial advantage and that copies bear this notice and the full citation on the first page. Copyrights for third-party components of this work must be honored. For all other uses, contact the owner/author(s). Gen-IR@SIGIR2024, July 18, 2024, Washington DC, US © 2024 Copyright held by the owner/author(s). retrieval for complex queries. Dense retrieval has gained popularity due to transformer-based models that map queries and documents into continuous, high-dimensional vector spaces using a transformer-based encoder (BERT, RoBERTa), enabling better representation of semantic meanings. This approach addresses some limitations of sparse retrieval, such as the lack of semantic similarity and vocabulary mismatches. Despite its advantages, dense retrieval presents challenges, including fine-grained interactions between passages and queries [24] , embedding space bottleneck [11] , and the need for large pre-computed indexes that demand substantial memory and resources [27] . Additionally, these models lack end-to-end processing capabilities.To address the limitations previously discussed, a series of works [4, 5, 22] have proposed alternative approaches in generative retrieval. These end-to-end models utilize sequence-to-sequence language models, such as T5 [19] and BART [12] , to retrieve documents by generating their unique document identifier (docID). This method supports fine-grained interaction between queries and do-cIDs, allowing for comprehensive end-to-end optimization [25] . However, most of these approaches overlook the interaction between the query and the context of its relevant document, where this interaction occurs only within the parameters and at the token prediction stage. This oversight means that the model does not adequately capture the relationship between the context of a document and the query. Consequently, the model may only thoroughly learn some nuanced interactions between them, potentially affecting retrieval performance.In order to achieve an improved retrieval performance, we propose a Document Aware Generative Retriever (DAGR) model leveraging a novel document-aware learning technique that enhances query representations to better align with document representations. This is accomplished by incorporating a margin-based contrastive loss at the encoder's output stage. We also establish a connection between dense retrieval and the language model head by introducing a method known as document-aware Label Mapping (DALM). DALM leverages the principles of dense retrieval, applying contrastive loss to the document embedding and the language model head's weights associated with the document's respective tokens. Furthermore, we present a curriculum-based learning approach that incrementally increases the difficulty of negative samples used in contrastive loss, fostering a more robust learning process. Additionally, we explore the application of contrastive loss at the token level to further enhance model performance.This study aims to address several research questions:• RQ1: Investigate the effectiveness of document-aware curriculum learning and label mapping. • RQ2: Evaluate the impact of curriculum learning on the performance of contrastive loss.• RQ3: Compare the efficacy of margin loss and listwise loss.• RQ4: Assess the performance of different models on documents not encountered during post-tuning. • RQ5: Examine the advantages of DAGR over token-level contrastive learning. 
related work 2.1 information retrieval: 
 Retrieval is a fundamental part of many Knowledge Intensive NLP tasks. The task involves retrieving relevant information from a large corpus of knowledge based on a given query. Formally, given a query q and a corpus C, the goal is to retrieve the top k documents from C that are most relevant to q, where relevance is typically measured using a similarity metric between the query and documents in the corpus.2.1.1 Sparse Retrieval. Sparse Retrievers are a retrieval method that uses explicit lexical matches between the query and the documents in the corpus. In other words, they retrieve documents that contain specific keywords or phrases that match the query. One common approach is to use an inverted index, a data structure that stores all the terms in the corpus along with the list of documents in which they appear. Given a query, the inverted index can quickly retrieve the documents that contain the query terms. Some examples of sparse retrieval methods include TF-IDF and BM25.2.1.2 Dense Retrieval. Dense retrieval is a method that leverages transformer-based models to encode queries and documents from a corpus into dense vectors. These methods retrieve documents that closely resemble the query vector, capturing semantic similarity between the query and documents. Unlike sparse retrieval, dense retrieval can find relevant documents even when they do not contain exact query term matches. Dense retrieval methods often involve pre-computed indexing, which converts documents to vectors and stores them for later access during searches. Notable examples include dual-encoder dense passage retrieval techniques [9, 20] . Dense retrieval uses dense vector representations of text for similarity search. In this approach, a neural network encodes the query and the document as dense vectors. The similarity between the two vectors is then computed using distance metrics such as cosine similarity, allowing dense retrievers to capture the semantics of the text for improved retrieval performance over sparse retrievers. This method often employs pre-trained language models like BERT and RoBERTa, which can capture rich and complex semantic relationships for enhanced retrieval tasks. While dense retrieval offers several advantages, such as semantic understanding and improved relevance, it also presents challenges. Computing dense vector representations can be computationally intensive and slow, mainly when dealing with a large number of documents.2.1.3 Generative Retrieval. Generative retrieval is an end-to-end retrieval technique that generates text sequences as a response to a query rather than retrieving pre-existing text embedding. In this approach, initially, every document is mapped to a unique identifier that represents a passage or document like Passage title [5] , substrings [4] , URL [27] pseudo queries [14] and sometimes the whole passage [11] . More recent models have explored semantic docid, which is constructed using either hierarchical K-means [22, 24] or based on product quantisation [27] . These models typically rely on sequence-to-sequence transformer architectures [23] built on top of models like T5 [19] , BART [12] . They are trained to predict the sequence of tokens in a docid given an input query. The objective of the retrieval can interpreted asP (𝑑 |𝑞) = |𝑑 | 𝑖=0 P (𝑑 𝑖 |𝑞, 𝑑 <𝑖 )for a document with docid 𝑑 and it's respective query 𝑞, and 𝑑 𝑖 is the 𝑖th token of the docid 𝑑. During inference, these models use beam search to retrieve or generate ranked lists of documents. To mitigate hallucinations and avoid generating unwanted tokens in the docID, constrained beam search is employed, guiding the language model to generate text that follows certain constraints using data structures such as prefix trees or FM indexes. 
methodology: document aware training for gr: 
 In this work, we propose a two-stage training procedure for generative retrieval. Initially, we train a straightforward sequenceto-sequence transformer [23] based on the T5 architecture [19] . This model processes a query as input and produces the most relevant document by generating a sequence of tokens that form the document's identifier (docID). The training utilizes a substantial collection of (query, docID) pairs, allowing the model to map queries to document identifiers effectively. Following recent works [22, 24] , we employ a hierarchical k-means algorithm to construct the do-cIDs. This clustering approach organizes documents hierarchically, facilitating the generation of meaningful docIDs that capture the semantic structure of the corpus. In the second stage, we fine-tune the trained model using a novel document-aware curriculum-based contrastive learning approach, which further enhances the model's performance. Inspired by the dense passage retrieval (DPR) technique [9] , this method applies margin-based/listwise contrastive loss to minimize the distance between the encoder/decoder representation of the query and the corresponding document, aiming to achieve similar representations for both. Additionally, we leverage the document embedding to create document-aware label mappings, reducing the distance between the document embedding and the language model head's weights associated with the respective tokens. This approach enables the model to align document embeddings with the output space better, leading to improved retrieval performance.In our approach, Similar to NCI [24] , we create document identifiers based on the input query while integrating document semantic information into the retrieval process. This technique helps manage large corpora by embedding useful semantic content into the identifiers, allowing us to encode the documents' semantics into the learning process. The goal is to generate document identifiers that are semantically aligned so documents with similar content have closely related identifiers, thereby streamlining the retrieval process. We employ the hierarchical k-means algorithm to cluster and encode documents to accomplish this. This process begins by classifying the documents into 𝑘 clusters using their BERT-encoded [6] representations. The algorithm is applied recursively for clusters with more than 𝑐 documents. Once each cluster is limited to 𝑐 documents or fewer, the documents within each cluster are numbered sequentially from 0 to 𝑐 -1. This method organizes the documents into a tree structure (𝑇 ) with a root node (𝑟 0 ). Each document is linked to a leaf node through a path (𝑙 = 𝑟 0, 𝑟 1, ..., 𝑟𝑚) from the root, where each internal node represents a cluster at different levels of the tree, and the leaf node corresponds to a specific document. The final identifier for each document is constructed by concatenating the indices from root to leaf, forming a unique semantic identifier for each document. Encoding documents this way ensures that documents with similar semantics have similar identifiers, making the retrieval process more efficient and accurate. Following the implementation of NCI, in our experiments, we use k = 30 and c = 30 as our parameters for clustering.Generating semantic document identifiers (docid) based solely on a single input query may lead to an inadequate representation of document semantics. Therefore, following the approaches used in NCI [24] and DSI [22] , we integrate each document's content into the model's parameters during training to enhance semantic capture. To achieve this, we augment queries using two methods: "document as queries" and "generated queries" from DocT5Query [16] . For the "document as queries" approach, we use the first 64 terms from a document as queries and randomly sample 10 additional groups of 64 consecutive terms from the entire document. This provides a more comprehensive representation of the document's content and semantics. For the "generated queries" method, we employ a sequence-to-sequence transformer [23] based on the DocT5Query architecture. This model takes the document tokens as input and generates relevant queries through random sampling.Random sampling ensures diversity in the generated queries and captures different aspects of the document's content.The base retriever employs a T5-based architecture [18] and adheres to a traditional learning objective for generative retrieval. This model processes an input query and learns to predict the sequence of tokens representing the relevant document's docid. Given a query 𝑞, the model's encoder produces a representation that is used by the decoder to generate the docid 𝑟 = 𝑟 EQUATIONwhere 𝜃 represents the model parameters and 𝑑 <𝑘 denotes the sequence of tokens preceding 𝑑 𝑘 . The probability function is defined as:EQUATIONwhere 𝑊 𝑙𝑚 ∈ R ℎ×𝑣 denotes the weights of the language modeling head, with ℎ being the model's hidden size and 𝑣 the decoder's vocabulary size. The query 𝑞 encompasses all ground truth queries, generated queries, and documents treated as queries. The training process optimizes the model parameters 𝜃 jointly across these various types of queries.Based on the model produced by the above process, we predict a few of the most relevant documents for each query. We accumulate all retrieved documents along with the relevant documents to build a dictionary {𝑑 : 𝑠}, where 𝑑 is the docid and 𝑠 is the set of most relevant documents according to the model. This set 𝑠 is a source of soft negatives when sampled from it. The document-aware loss utilizes the decoder output of the query and the document context 𝑑𝑐, whose dot product is calculated at the respective positions and used to build a listwise loss. This 𝑑𝑐 comprises 64 consecutive terms from the document, ensuring maximal overlap with the query 𝑞. The decoder query representation at 𝑖th position of the docid is calculated as follows:EQUATIONWhere D 𝑖 𝑞 is the query's decoder representation at the 𝑖th step. The similarity score of query decoder embedding D 𝑖 𝑞 and the decoder embedding of the positive document context (D 𝑖 𝑑𝑐 𝑝𝑜𝑠) is calculated and summed for all the position as shown below:EQUATIONThe scores for soft Negatives (S 𝑑 𝑛𝑒𝑔 ) from the most relevant documents according to the base model and the random negatives (S 𝑑 𝑛𝑒𝑔 ), which is taken from within the batch, is calculated in a similar fashion and the decoder document aware loss is calculated as follows:EQUATIONWe also utilize the average pooled output of the Encoder to calculate a margin loss between the query and its corresponding document context 𝑑𝑐.The average pooled output of the Encoder is calculated as follows:EQUATIONwhere E 𝑞 ∈ R 𝑏 ×ℎ represents the encoder output. The distances between the query embedding E𝑞 and the embeddings of the positive document context (E 𝑑𝑐 𝑝𝑜𝑠 ), soft negatives (E 𝑑𝑐 𝑛𝑒𝑔 ), and random negatives (E 𝑑𝑐 𝑟𝑛𝑑 ) are calculated as follows:EQUATIONEQUATIONEQUATIONThe margin loss can be calculated using the following equation:EQUATION)where 𝛾 is a predefined margin that controls the separation between positive and negative samples.This margin loss helps the model learn to differentiate between the query and relevant and non-relevant documents, thereby improving retrieval performance. The overall document-aware loss is calculated as follows:L 𝑒𝑛𝑐 𝑑𝑎 = 0.5 • L 𝑟𝑛𝑑 + 0.5 • L 𝑛𝑒𝑔 (12)Easy to Hard Curriculum learning [2] is a training strategy in which the complexity of the training samples is gradually increased over time to improve the model's learning process. In the context of contrastive learning for document retrieval, we apply curriculum learning to sample negatives of varying difficulty as training progresses. This approach is closely related to the curriculum knowledge distillation done in GripRank [1] . The curriculum learning approach can be described as follows:(1) Initial Steps (0 ≤ 𝑡 ≤ 𝑇 0 ) In the initial steps of training, the model randomly samples from the top 100 documents retrieved by the model as soft negatives. This introduces diversity in the negative samples and helps the model learn general patterns in the data. (2) Intermediate Steps (𝑇 0 < 𝑡 ≤ 𝑇 1 ): As training progresses, the model begins sampling from the top 5 documents with a probability calculated as:EQUATION)where 𝑡 represents the current training step and 𝑇 1 is the end of the intermediate training phase. This gradual transition increases the difficulty of the negative samples, encouraging the model to focus on more challenging cases. (3) Later Steps (𝑡 > 𝑇 1 ): After 𝑇 1 , the model samples from the most relevant documents as the soft negatives for the remaining training steps with a probability of around 0.75. This exposes the model to increasingly difficult negatives, helping it fine-tune its understanding of the most relevant documents. By gradually increasing the difficulty of the training samples using curriculum learning, the model can better focus on differentiating between relevant and non-relevant documents, thereby enhancing its retrieval performance.The Language Model weights 𝑊 𝑙𝑚 at the top of the decoder can be viewed as a scope-reduced document index, drawing similarities to techniques like product quantization. The decoder's output at the 𝑖-th position serves as a query representation, while the 𝑊 𝑙𝑚 at the 𝑖-th position consists of 𝑐(number of child per node) relevant vectors. Essentially, the model aims to find the vector with the highest dot product, corresponding to the respective token in the docid. This process resembles a dense retriever with scope reduction algorithms such as product quantization.Building on this concept, we aim to maximise the similarity between the weights of a token corresponding to the docid (𝑊 𝑙𝑚 EQUATIONSimilarly the similarity score S 𝑙𝑚 𝑝𝑜𝑠 and S 𝑙𝑚 𝑝𝑜𝑠 is calculate for the negative and random passage token representations D 𝑖 EQUATIONAnd the main model is trained by optimising all the loss joints as mention in the following equation:EQUATION4 EXPERIMENTAL SETUP 4.1 Dataset 4.1.1 Natural Questions. We conduct training for the retrieval model using a widely recognized open-source question answering dataset called Natural Questions (NQ) [10] . Introduced by Google in 2019, Natural Questions consists of approximately 320,000 querydocument pairs, providing a substantial dataset for training and evaluating retrieval models. The documents in the dataset are sourced from Wikipedia pages, offering a broad and diverse range of topics and domains. The queries in NQ are naturally phrased questions, mirroring real-world information-seeking behavior and presenting a practical challenge for retrieval systems. This diversity and realism in the dataset allow models to generalize better to a variety of queries and document contexts.  
implementation details: 
 For the hierarchical semantic identifier, the k-means algorithm is applied to the document embeddings obtained from a 12-layer BERT model. The number of centroids (𝑘) at each layer is set to 30, and the maximum number of clusters for the terminal condition is set to 30 as well. For training the baseline model, we use the T5 implementation from the Hugging Face Transformers library. We initialize the weights of the encoder using the pre-trained T5 model [18] , while the decoder weights and other components are randomly initialized. The initial base model is trained with a learning rate of 1 × 10 -4 on two NVIDIA V100 GPUs 32 GB of memory each, and a batch size of 128 per device, resulting in a total batch size of 256. The maximum token length for all inputs is set to 64. For the document-aware model, we train on two NVIDIA V100 GPUs with a batch size of 32 per device and a gradient accumulation of 4.The margin value (𝜆) is set to 1 for all margin loss calculations, and the token size for both the query and document context is set to 64.Regarding curriculum learning, we define 𝑇 0 as the number of steps in one epoch and set 𝑇 1 as the total number of steps in 15 epochs. Base model is trained on the dataset with queries from ground truth queries, pseudo queries from DocTTTTTQuery and document as queries for 850k steps. For the DAGR model we train on the top base model with L 𝑑𝑎𝑔𝑟 on only the ground truth queries for 50k steps. We follow the same steps provided in the NCI's [24] GitHub page. The NQ dataset consist of around 100k unique documents and around 300k unique queries.We compare our models to strong baselines across all types of retrieval models. For sparse retrieval, we compare our models to BM25 using raw documents and an approach using augmented queries from DocT5Query [16] . For dense retrieval models, we compare our models with DPR [9] and ANCE [26] . For generative retrieval models, we use DSI [22] , GENRE [5] , SEAL [4] , Ultron [27] , and NCI [24] as baselines for comparison. All the results are sourced either from [24] or their respective original papers.Recall@n is the standard evaluation metric for the retrieval task, this is formulated as follows:EQUATIONIn this case |(relevant article)| = 1 for all the query will makes 𝑅𝑒𝑐𝑎𝑙𝑙 = |(relevant article) ∩ (retrieved article)|, and this recall will be "1" if the model retrieves the relevant article and "0" otherwise. R@k accounts for total article retrieved by the model which k, it is controlled by user. Recall in this case can also be consider as accuracy as for R@k the model will retrieve k articles and we score 1 if the relevant article is present in the k article retrieved and "0" other wise this similar to the accuracy metric.Mean Reciprocal Rank (MRR) is a metric used to evaluate the effectiveness of search engines, recommendation systems, and information retrieval systems. It is particularly common in the context of ranking-based evaluations. MRR is defined as the average of the reciprocals of the ranks of the first correct item.EQUATIONWhere:• |𝑄 | the total number of queries or instances.• 𝑟𝑎𝑛𝑘 𝑖 is the rank of the first correct item for the 𝑖-th query. 
retrieval performance on nq: 
 From table 1 our model DAGR performs modestly better compared to most of the models from respective class of retrievers across all 4 different metrics on Natural questions. We can notice that NCI is second best performing according to R@1 and MRR@100. DAGR also out performs our Base T5 model suggesting the importance of the document aware representations. The table 2 show the result of ablation studies. we do experiments by removing L 𝐷𝐴𝐿𝑀 or curriculum and study the performance of the model. we also replace the L 𝑒𝑛𝑐 𝑐𝑎 with Listwise loss similar to L 𝑑𝑒𝑐 𝑑𝑎 and compare the performance to the token based learning to rank. Without DALM Loss (RQ1): Excluding the L 𝐷𝐴𝐿𝑀 loss from the model training results in a performance decrease of approximately 0.36%. This suggests that incorporating the document-aware loss to modify the language model head is beneficial for the model to learn token weights that are closer to the embeddings related to the context of the document.Without Curriculum Learning (RQ2): When not using curriculum learning and instead using only the top 5 most relevant documents as negative samples, performance decreases by around 0.33 in Recall@1%. Figures 1 illustrate that contrastive loss is lower when curriculum learning is used compared to not using it. Furthermore, even several steps after the curriculum phase, the loss remains consistently lower. This attributes to the efficacy of curriculum learning to allow model some time learn better representations in initial steps and lets it do well for harder negatives in later stages Listwise Loss (RQ3) and L 𝑙𝑡𝑟 (Learning to Rank) Loss (RQ5): Comparing the model performance using listwise and LTR loss, the DAGR model achieves improvements of 0.6% and 0.9% in Recall@1 metric respectively. Performance of margin loss at encoder level is also better than both the Listwise loss and L 𝑙𝑡𝑟 .Performance on Unseen Documents While Post-tuning (RQ4): We evaluate the model's performance on documents that were not present during post-tuning. These documents don't have ground truth queries but are included in the test set, allowing us to assess the model's ability to generalize and avoid overfitting. As shown in Table 3 , the zero-shot performance of DAGR is slightly better than that of the baseline T5 model, suggesting that the model's generalization capability is on par with or slightly better than the base model. This indicates that our approach effectively balances training on the available data while still maintaining strong performance on unseen documents. 
conclusion: 
 This work delves into the one possible way to interact query and document during the training. Since this interactive training is done after the initial generative retrieval model any model could be further fine-tuned with this technique. The experiment on using this fine-tuning approach on models with PAWA [24] and other generative retrieval models could be considered for a future work. 
 
 Appendix:  


Decision: ## Decision: Accept

The authors requested reviews to remain private.