CODER : An efficient framework for improving retrieval through CO ntextual Document Embedding Reranking (2022)

Abstract: Contrastive learning has been the dominant approach to training dense retrieval models. In this work, we investigate the impact of ranking context - an often overlooked aspect of learning dense retrieval models. In particular, we examine the effect of its constituent parts: jointly scoring a large number of negatives per query, using retrieved (query-specific) instead of random negatives, and a fully list-wise loss.To incorporate these factors into training, we introduce Contextual Document Embedding Reranking (CODER), a highly efficient retrieval framework. When reranking, it incurs only a negligible computational overhead on top of a first-stage method at run time (approx. 5 ms delay per query), allowing it to be easily combined with any state-of-the-art dual encoder method. Models trained through CODER can also be used as stand-alone retrievers.Evaluating CODER in a large set of experiments on the MS MARCO and TripClick collections, we show that the contextual reranking of precomputed document embeddings leads to a significant improvement in retrieval performance. This improvement becomes even more pronounced when more relevance information per query is available, shown in the TripClick collection, where we establish new state-of-the-art results by a large margin.

## introduction neural text retrieval models typically rely on a contrastive training optimization that uses (query, positive document, negative document) triplets as training samples. this scheme is especially popular, as it is well-suited to the computational constraints of large transformer-based language models such as bert @xcite and its variants for retrieval @xcite @xcite . referred to as pair-wise training, the model is asked to score the similarity between the query embedding and a ground truth relevant (positive) document embed-ding, higher than the one between the query and a negative document embedding. while such contrastive learning approaches are effective, by only considering pairs of positive and negative documents at a time, they (1) deviate from the target objective of comparing a query against many documents while discarding inter-document information, and (2) depart from core list-wise evaluation metrics like ndcg @xcite . addressing the first shortcoming, recent works have employed "in-batch" negatives: given a batch containing (query, positive document) tuples, the negatives to a tuple are set as the known positive documents from other queries within that batch. (e.g. @xcite ; @xcite ; @xcite ; @xcite ; @xcite ). although this approach efficently increases the number of negatives to improve performance, the presence of hard negative samples foot_0 is critical in achieving state-of-the-art (sota) @xcite @xcite . the rich literature on learning-to-rank (l2r) has outlined compelling reasons for taking into account the context of other candidate documents being ranked for a query when scoring each document @xcite @xcite , realized in various list-wise optimizations. such approaches allow for the model to directly optimize ir metrics as opposed to a surrogate pair-wise loss as seen in the contrastive training regime. despite these list-wise approaches achieving competitive results in a variety of ranking situations, considerations of computational complexity and stochastic stability practically relegate them to shallow neural models over handcrafted feature vectors @xcite @xcite . in this paper, we extend existing work in negative sampling and list-wise learning to allow for large pre-trained language models to take advantage of the context found in a large, coherent set of candidate documents. we particularly examine the effect of constituent parts of the query context, i. e. (1) jointly scoring a large number of negatives, (2) using retrieved (query-specific) instead of random negatives, and (3) a fully list-wise loss. to this end, we introduce contextual document embedding reranking (coder), a highly efficient and generic fine-tuning framework that for the first time enables incorporating context, previously only considered in learning-to-rank neural networks, into transformer-based language models used in state-of-the-art dense retrieval. coder acts as a lightweight performance enhancing framework that operates on precomputed document embeddings, while transforming the query to account for new list-wise context information over a large number of query-specific hard negative candidate documents. it can be applied to virtually any existing dual-encoder model, and used both for singleas well as two-stage dense retrieval. our contribution is three-fold: (1) we introduce an efficient framework which enables leveraging ranking context; (2) we conduct a large set of experiments on the ms @xcite and tripclick @xcite collections and show that coder can considerably enhance the effectiveness of a wide class of dense retrieval models at minimal computational cost, while achieving new sota results on tripclick; ## related work recent work has demonstrated the importance of the quality of negative documents used during finetuning. @xcite periodically re-encode every query and document in the collection during training in order to mine the most difficult documents to use as negative candidates via approximate nearest neighbor (ann) search. improving on this slow and resource-intense process, @xcite (published as @xcite ) forego fine-tuning of the document encoder, instead only fine-tuning the query encoder while dynamically mining negatives. tas-b @xcite also improves the quality of negatives by clustering semantically similar queries, such that the in-batch negatives are indirectly related to the ground truth document. while we contribute to this line of research, we show that one can avoid such complexity and directly benefit from list-wise optimization applied on a large, coherent and informative context of candidate documents, retrieved for each query in advance. moving from quality to quantity of negative candidates, rocketqa @xcite drastically increases the quantity of random in-batch negatives to several thousands by sharing negatives across at least 8 v100 gpu instances. despite its huge size, this pool of sampled negatives includes only 4 retrieved hard negatives per query and is otherwise almost entirely random, with no shared context across documents. unlike coder, this approach necessitates training an expensive cross-encoder model to "denoise" (filter out) retrieved candidates, otherwise yielding poor performance. moreover, it involves using the crossencoder to pseudo-label additional data samples, an approach adopted by the currently top performing dual encoder models following up this work, rocketqav2 @xcite , which additionally leverages list-wise cross-encoder teaching, and pair @xcite , which includes a loss term capturing similarity between passages. to address the increasing complexity and computational requirements of training pipelines, @xcite instead propose corpusspecific, self-supervised pre-training with a bespoke transformer add-on (see baselines in section 4). an orthogonal approach to reduce computational requirements focuses on jointly optimizing query optimizer and product quantization for ann search @xcite . list-wise loss functions have been extensively used within learning-to-rank (l2r), although they have been limited to either shallow neural models @xcite or various deep networks @xcite ; @xcite ; @xcite in works focusing on l2r datasets. these consist of handcrafted feature vectors representing query-document similarity such as term overlap, click-through rate, bm25 scores, and other salient features. effectively applying l2r concepts to transformer-based language models used for ad-hoc dense retrieval is a nontrivial challenge, and represents coder's exten-sion of prior works. finally, there is a body of work utilizing large pre-trained language models for retrieval in crossencoder @xcite @xcite , late cross-encoder @xcite , and generative rankers @xcite , as well as query/document expansion and indexing @xcite @xcite @xcite . co-bert @xcite , a recent method leveraging ranking context, uses a cross-encoder bert reranker to select candidates and to compute feature vectors as input for the l2r methods above @xcite through query-document term interactions. in comparison, coder is orders of magnitude more efficient both during training and inference. all existing approaches either advocate for using a handful of hard negatives (e.g. @xcite ), or compromise with it due to computational constraints. by recognizing the importance of context, our proposed framework is the first to allow the combination of quantity and quality of negatives for training sota dual encoder models with very modest computational resources. ## method coder involves fine-tuning a pre-trained query encoder to learn a query representation that is as proximal as possible to the representation of the ground-truth relevant document(s), by adjusting it to better account for the context of multiple queryrelated documents. the architecture consists of two main components (fig. 1 ): a query encoder, which builds a query representation, and the document set scoring module, which, given a query representation, jointly scores a set of n precomputed embeddings of positives and hard negatives retrieved by an arbitrary retrieval method, m c . using precomputed document embeddings reduces computational costs (memory, flops) by a factor of n . thus, unlike all existing approaches, we can afford to use a large number of such hard negatives (@xmath0, unless otherwise noted). ## experimental setup datasets. we conduct the experiments on passage and document retrieval tasks, foot_1 using two large publicly available ir collections in the domains of web and health retrieval. the first dataset is ms marco passage retrieval dataset @xcite used for training and evaluation. we evaluate the models trained on ms marco also on trec deep learning passage retrieval track 2019 and 2020 @xcite . the second dataset is tripclick, a recently introduced health document retrieval dataset @xcite used for training and evaluation. details of the collections are provided in section a.2 of the appendix. while the training data of ms marco contains only approximately 1 relevance judgement per query, the tripclick collection has the advantage of providing a much larger set of relevance information, namely approximately 42, 9, and 3 data points per query in head, torso, and tail sets, respectively. as shown in the next section, this is particularly beneficial when optimizing over a large ranking context in a list-wise manner. evaluation details are given in section a.3 of the appendix. baselines. we choose several dense retrieval models as baselines, i. e. "base models" subjected to coder fine-tuning: 1. repbert @xcite , a bert-based model with a typical dual encoder architecture which underpins all state-of-the-art dense retrieval methods, trained using a triplet max-margin loss. 2. tas-b @xcite , which, besides being a top-performing dense retrieval method on the ms marco / trec-dl 2019, 2020 datasets, it also represents methods that have been optimized with respect to their training process (details in section 2). 3. finally, to explore the limits of coder, we use it to fine-tune a trained cocondenser retriever @xcite , the state-of-the-art dense retrieval model that does not make use of querydocument term interactions, cross-encoder teacher models or additional pseudo-labeled data samples, but instead relies on extensive corpus-specific, selfsupervised pre-training using a special architecture and contrastive loss component. it is a particularly challenging baseline for our coder framework, because it has been trained through (a) mining for hard negatives using a trained version of the model itself, and (b) the negative loglikelihood (infonce) loss, which is "nearly" list-wise (it differs from our kl-divergence loss only when there are more than one positive candidates). configurations. for the coder framework, we use the following notation: m @xmath1, m c ), where m d is the base model used to encode documents into document embedding vectors (and initialize the query encoder weights), m c is the first-stage retrieval method used to procure the candidate (context) documents reranked during the coder training process, and m f is the retrieval method used as a first stage when coder is used as a reranking method during inference; m f vanishes in case coder is used directly for single-stage dense retrieval, and the notation coder(m d , m c ) is used instead. in addition to tas-b and repbert, we also experiment with bm25 (anserini implementation @xcite ) as m c and m f methods. specific hyperparameter details can be found in ## conclusion we examine the importance of ranking context and the effect of its constituent parts, i.e., a fully list-wise loss, a large number of negatives, and retrieved (query-specific) instead of random negatives, and show that they are all important ingredients for improving performance. we demonstrate that a lightweight reranking framework designed to leverage context is sufficient to significantly enhance the effectiveness of a wide range of dense retrieval models, without expensive cross-encoder distillation, pseudo-labeling or "denoising" negatives. the only computational overhead is a fast, resource-light fine-tuning process, with little (when the model is used as a reranker) to no (when used as a single-stage retriever) extra computational cost during inference.