Aggretriever: A Simple Approach to Aggregate Textual Representations for Robust Dense Passage Retrieval (2023)

Abstract: Pre-trained language models have been successful in many knowledge-intensive NLP tasks. However, recent work has shown that models such as BERT are not “structurally ready” to aggregate textual information into a [CLS] vector for dense passage retrieval (DPR). This “lack of readiness” results from the gap between language model pre-training and DPR fine-tuning. Previous solutions call for computationally expensive techniques such as hard negative mining, cross-encoder distillation, and further pre-training to learn a robust DPR model. In this work, we instead propose to fully exploit knowledge in a pre-trained language model for DPR by aggregating the contextualized token embeddings into a dense vector, which we call agg★. By concatenating vectors from the [CLS] token and agg★, our Aggretriever model substantially improves the effectiveness of dense retrieval models on both in-domain and zero-shot evaluations without introducing substantial training overhead. Code is available at https://github.com/castorini/dhr .

## introduction a bi-encoder architecture @xcite based on pre-trained language models @xcite @xcite has been widely used for first-stage retrieval in knowledge-intensive tasks such as open-domain question answering and fact checking. compared to bag-of-words models such as bm25, these approaches circumvent lexical mismatches between queries and passages by encoding text into dense vectors. despite their success, recent research calls into question the robustness of these singlevector models @xcite . as shown in figure 1 , single-vector dense retrievers (e.g., bert cls and tas-b) trained with well-designed knowledge distillation strategies @xcite still underperform bm25 on out-of-domain datasets. along the same lines, @xcite find that simple entity-centric questions are challenging to these dense retrievers. recently, @xcite observe that pre-trained language models such as bert are not ''structurally ready'' for fine-tuning on downstream retrieval tasks. this is because the [cls] token, pre-trained on the task of next sentence prediction (nsp), does not have the proper attention structure to aggregate fine-grained textual information. to address this issue, the authors propose to further pre-train the [cls] vector before fine-tuning and show that the gap between pre-training and fine-tuning tasks can be mitigated (see cocondenser cls illustrated in figure 1 ). however, further pre-training introduces additional computational costs, which motivates us to ask the following question: can we directly bridge the gap between pre-training and fine-tuning without any further pre-training? before diving into our proposed solution, we briefly overview the language modeling pretraining and dpr fine-tuning tasks using bert. figure 2 (a) illustrates the bert pre-training tasks, nsp and mask language modeling (mlm), while figure 2 (b) shows the task of fine-tuning a dense retriever. we observe that solely relying on the [cls] vector as the dense representation does not exploit the full capacity of the pre-trained model, as the [cls] vector participates directly only in nsp during pre-training, and therefore lacks information captured in the contextualized token embeddings. a simple solution is to aggregate the token embeddings by pooling (max or mean) into a single vector. however, information is lost in this process and empirical results do not show any consistent effectiveness gains. hence, we see the need for better aggregation schemes. in this paper, we propose a novel approach to generate textual representations for retrieval that fully exploit contextualized token embeddings from bert, shown in figure 2(c ). specifically, we reuse the pre-trained mlm head to map each contextualized token embedding into a highdimensional wordpiece lexical space. following a simple max-pooling and pruning strategy, we obtain a compact lexical vector that we call agg . by concatenating agg and the [cls] vector, our novel aggretriever dense retrieval model captures representations pre-trained from both nsp and mlm, improving retrieval effectiveness by a noticeable margin compared to fine-tuned models that solely rely on the [cls] vector (see bert agg vs bert cls in figure 1 ). importantly, fine-tuning aggretriever does not require any sophisticated and computationally expensive techniques, making it a simple yet competitive baseline for dense retrieval. however, our approach is orthogonal to previously proposed further pre-training strategies, and can still benefit from them to improve retrieval effectiveness even more (see cocondenser agg in figure 1 ). to the best of our knowledge, this is the first work in the dpr literature that leverages the bert pre-trained mlm head to encode textual information into a single dense vector. ## background and motivation given a query q, our task is to retrieve a list of passages to maximize some ranking metric such as ndcg or mrr. dense retrievers @xcite based on pre-trained language models encode queries and passages as low dimensional vectors with a bi-encoder architecture and use the dot product between the encoded vectors as the similarity score: recent work @xcite shows that the [cls] vector remains ''dormant'' in most layers of pre-trained models and fails to adequately aggregate information from the input sequence during pre-training. thus, researchers argue that the models are not ''structurally ready'' for fine-tuning. to tackle this issue, unsupervised contrastive learning has been proposed, which creates pseudo relevance labels from the target corpus to ''prepare'' the [cls] vector for retrieval. the most representative technique is the inverse cloze task (ict; @xcite . however, since the generated relevance data is noisy, further pre-training with ict often requires a huge amount of computation due to the need for large batch sizes or other sophisticated training techniques @xcite @xcite . another thread of work @xcite manages to guide transformers to aggregate textual information into the [cls] vector through auto-encoding. this method does not require as much computation as unsupervised contrastive learning but is still much more computationally intensive than fine-tuning. for example, @xcite report that the further pre-training process still requires one week on four rtx 2080 ti gpus, while fine-tuning consumes less than one day in the same environment. recent work on neural sparse retrievers @xcite projects contextualized token embeddings into a highdimensional wordpiece lexical space through the bert pre-trained mlm projector and directly performs retrieval in wordpiece lexical space. these models demonstrate that mlm pre-trained weights can be used to learn effective lexical representations for retrieval tasks, a finding that has not been fully explored in the dpr literature. inspired by this work, we explore reusing mlm pre-trained weights for dpr fine-tuning and further combine the [cls] vector to fully exploit textual information in a pre-trained language model. ## aggretriever in this section, we first introduce our method for text aggregation to form agg , which consists of two steps: pooling and pruning. then, we describe how to concatenate the aggregated text representation agg and [cls] into a 768-dimensional dense vector for fine-tuning and retrieval. ## related work dense retrieval. the most related line of research to our own work is the literature on how to effectively fine-tune a single-vector dense retriever. on the one hand, some researchers propose computationally expensive fine-tuning techniques such as hard negative mining strategies @xcite , knowledge distillation @xcite , or their combination @xcite . on the other hand, others leverage further pre-training to improve the subsequent fine-tuning @xcite @xcite @xcite @xcite . as far as we are aware, our work is the first to discuss how to fine-tune dense retrieval models to effectively aggregate textual information from the pre-trained mlm head rather than directly using the [cls] vector or contextualized embeddings from max or average pooling @xcite . sparse retrieval. previous work @xcite @xcite has demonstrated that projecting contextualized token embeddings into a high-dimensional vector in the wordpiece vocabulary space is an effective way to represent tokenlevel information from transformers for lexical matching. these models directly feed the highdimensional vectors into an inverted index for retrieval. thus, sparsity control for effectivenessefficiency tradeoffs involves additional considerations @xcite . in contrast, our approach converts high-dimensional vectors into low-dimensional ones where top-k retrieval can be performed directly using ann search libraries @xcite . hybrid retrieval. our work can be characterized as hybrid since we ''fuse'' semantic and lexical representations into a single dense vector. recent work @xcite @xcite proposes to jointly train [cls] and token-level representations for semantic and lexical matching, respectively. the two kinds of representations require different implementations for top-k retrieval, so multiple software stacks are required to perform retrieval. in contrast, our representations retain the best of semantic and lexical matching, but entirely as dense vectors. thus, retrieval can be performed in a simple execution environment. ## conclusion and future work in this paper, we present aggretriever, a singlevector dense retrieval model that exploits all contextualized token embeddings from the input to bert. we introduce a simple approach to aggregate the contextualized token embeddings into a dense vector, agg . experiments show that agg combined with the standard [cls] vector achieves better retrieval effectiveness than using the [cls] vector alone for both in-domain and zero-shot evaluations. our work demonstrates that mlm pre-trained transformers can be fine-tuned into effective dense retrievers without further pre-training or expensive fine-tuning strategies. our work leads to a few open questions for future research: (1) since we have demonstrated that aggretriever still benefits from further pretraining, can we design additional pre-training tasks tailored directly to our model? the design of these tasks, of course, needs to be mindful of the computational costs. (2) can we apply current state-of-the-art compression techniques to aggretriever? @xcite has shown that 768-dimensional dense representations can be effectively compressed into much smaller vectors. however, it is still unknown if these techniques can be applied to aggretriever to retain both in-domain and zero-shot retrieval effectiveness. (3) finally, can we apply aggretriever to multi-lingual retrieval? because, in a multi-lingual bert model, the mlm head can project into tokens in multiple languages, we can envision a natural extension. however, as shown in section 5.5, mlm projection is expensive, and the issue becomes worse when using a pre-trained multi-lingual model since the vocabulary size is usually even larger.