R etro MAE -2: Duplex Masked Auto-Encoder For Pre-Training Retrieval-Oriented Language Models (2023)

Abstract: To better support information retrieval tasks such as web search and open-domain question answering, growing effort is made to develop retrieval-oriented language models, e.g., RetroMAE and many others. Most of the existing works focus on improving the semantic representation capability for the contextualized embedding of the [CLS] token. However, recent study shows that the ordinary tokens besides [CLS] may provide extra information, which help to produce a better representation effect. As such, it’s necessary to extend the current methods where all contextualized embeddings can be jointly pre-trained for the retrieval tasks. In this work, we propose a novel pre-training method called Duplex Masked Auto-Encoder, a.k.a. DupMAE. It is designed to improve the quality of semantic representation where all contextualized embeddings of the pre-trained model can be leveraged. It takes advantage of two complementary auto-encoding tasks: one reconstructs the input sentence on top of the [CLS] embedding; the other one predicts the bag-of-words feature of the input sentence based on the ordinary tokens’ embeddings. The two tasks are jointly conducted to train a unified encoder, where the whole contextualized embeddings are aggregated in a compact way to produce the final semantic representation. DupMAE is simple but empirically competitive: it substantially improves the pre-trained model’s representation capability and transferability, where superior retrieval performances can be achieved on popular benchmarks, like MS MARCO and BEIR. We make our code publicly available at https://github.com/staoxiao/RetroMAE .

## introduction neural retrieval is important to many real-world scenarios, such as web search, question answer- †. equal contribution and designated as co-first authors. * . co-corresponding authors ing, and conversational system @xcite @xcite @xcite . in recent years, pre-trained language models, e.g., bert @xcite , roberta @xcite , t5 @xcite , are widely adopted as the retrievers' backbone networks. the generic pre-trained language models are not directly applicable to retrieval tasks. thus, it calls for complex fine-tuning strategies, such as sophisticated negative sampling @xcite , knowledge distillation @xcite , and the joint optimization of retriever and ranker @xcite . to reduce this effort and bring in better retrieval quality, there are growing interests in developing retrieval-oriented language models. one common practice is to leverage self-contrastive learning @xcite , where the language models are learned to discriminate heuristically acquired positive and negative samples in the embedding space. later on, auto-encoding is found to be more effective @xcite , where the language models are learned to reconstruct the input based on the generated embeddings. recent works @xcite further extend the auto-encoding methods by introducing sophisticated encoding and decoding mechanisms, which brings about remarkable improvements of retrieval quality on a wide variety of benchmarks. the existing retrieval-oriented pre-trained models mainly rely on the contextualized embedding from the special token, i.e., @xcite , to represent the semantic about input @xcite @xcite . however, recent study finds that other ordinary tokens may provide extra information and help to generate better semantic representations @xcite . such a statement is consistent with previous research @xcite , norwegian forest cat is a breed of domestic cat originating in northern europe encoder [cls] [m] forest cat is a breed of [m] cat originating in [m] europe cls embedding ot embeddings 1-layer trm cls embedding, masked input original input sentence lpu & aggregate ot embeddings bag-of-words (bow) figure 1: dupmae. encoder: the sentence is masked and encoded as the contextualized embeddings for [cls] and ordinary tokens. decoder: the cls embedding is joined with the masked input, where the original input is recovered by an 1layer transformer; ot embeddings are mapped into vocabulary space via lpu and aggregated to predict the bow feature. as multi-vector or token-granularity representations may give higher discriminative power than those based on one single vector. as a result, it is necessary to extend the previous works, such that the representation capability can be jointly pre-trained for both @xcite and ordinary tokens. to this end, we propose a novel auto-encoding framework called duplex masked auto-encoder, a.k.a. dupmae (figure 1 ). it employs two differentiated decoders working collaboratively, which aim to 1) improve each embedding's individual capacity, as well as 2) contribute to the quality of the joint representation derived from all embeddings. our proposed method is simple but empirically competitive. we perform dupmae on common pre-training corpus where a bert-based scale encoder is produced. our pre-trained model achieves superior performances in various downstream tasks. for supervised evaluations on ms marco, it reaches a mrr@10 of 42.6 in passage retrieval and a mrr@100 of 45.1 in document retrieval. for zero-shot evaluations on beir, it achieves an average ndcg@10 of 49.1 on all 18 datasets. it even notably outperforms strong baselines with more sophisticated fine-tuning approaches or much bigger model sizes. therefore, it validates that the representation capability and transferability of the pre-trained model can be substantially improved thanks to dupmae. ## related works neural retrieval is critical for many real-world applications, such as web search, question answering, advertising and recommender systems @xcite @xcite @xcite . it maps the query and document into embeddings within the same latent space, making their semantic relationship to be measured by the embedding similarity. in recent years, the pretrained language models have been widely applied to deep semantic retrieval such that discriminative representations can be generated for the queries and documents. despite the preliminary progress achieved by early pre-trained models, like bert @xcite , it is noticed that the more encoder cls 𝑥 ! 𝑥 " 𝒉# 𝑿 @xmath0 % + 𝐏 𝑥 & 𝑥 ! 𝑥 % 𝑥 " ⊕ p & p ' p ( p ) 𝒆 ' * p & p ' p ( p ) 𝑬 % + 𝐏 1 2 3 4 1 2 3 4 mask lpu 𝜇 ( 𝜇 ) 𝜇 ' 𝜇 max pooling bag-of-words (a) encoding (b) [cls] decoding (c) ot decoding 𝒆 ' + 𝒆 ' , 𝒆 ' - 𝒉 ' + 𝒉 ' , 𝒉 ' - 𝒉 ' + 𝒉 ' , 𝒉 ' - advanced models bring little benefit to the retrieval quality, and it's believed that the conventional pretraining algorithms are not compatible with the purpose of deep semantic retrieval @xcite @xcite . to mitigate the above problem, people become increasingly interested in developing retrieval oriented pre-trained models. for example, it is proposed to leverage self-contrastive learning (scl) where the language models are pre-trained to discriminate positive samples generated by data augmentation and in-batch negative samples @xcite @xcite . the scl based algorithms are limited by many factors, like the quality of data augmentation and the requirement of huge amounts of negative samples. later on, the auto-encoding based algorithms receive growing interests: the input sentences are encoded into embeddings, based on which the original sentences are reconstructed @xcite . the recently proposed methods, such as simlm @xcite and retro-mae @xcite , extend the previous autoencoding framework by upgrading the encoding and decoding mechanisms, which substantially improves the quality of deep semantic retrieval. the existing retrieval-oriented pre-training methods target on improving the semantic representation capacity for the contextualized embedding from the [cls] token. however, it is noticed that the ordinary tokens may provide additional information besides @xcite , especially when dealing with long and semantic-rich documents @xcite @xcite . as a result, it is necessary to extend the current works, where the representation capability can be enhanced for both types of contextualized embeddings. ## methodology we start with an overview of dupmae in this section. the framework of dupmae is shown as figure 2 . there is an unified encoder (a), where the masked input is encoded into its contextualized embeddings. there are two decoders working collaboratively. one decoder is applied for [cls] decoding (b): it employs a single-layer transformer, which reconstructs the original sentence based on the [cls] embedding. the other one is used for ot decoding (c): it utilizes a linear projection unit (lpu), which transforms the ordinary token embeddings into the vocabulary space. the transformed results are aggregated by max-pooling, where the bow feature of the input is predicted. the two decoding tasks are jointly conducted to train the encoder. the [cls] and ot embeddings are aggregated for the final representation of the input. with proper dimension reduction, it may preserve the same computation cost of inner-product and memory footprint as one single dense vector. ## experiment the empirical studies are conducted to explore the following research questions. benchmarks. the experiments are conducted for both supervised and zero-shot settings. we choose the passage and document retrieval task of ms marco benchmark @xcite for supervised evaluations. it contains queries from bing search, where ground-truth answers to the queries need to be retrieved from 8.8 million passages and 3 million documents, respectively. the queries from the dev set and trec deep learning track in 2019 (dl'19) @xcite are used for evaluation. we leverage beir benchmark @xcite for zero-shot evaluations. it contains a total of 18 datasets, which covers diverse types of retrieval tasks, such as question answering, duplication detection, and fact verification, etc. following the official evaluation script, the pre-trained models are fine-tuned with ms marco queries, and evaluated for their out-of-domain retrieval performances on each of the 18 datasets. baselines. we consider the following baselines for supervised evaluations according to their finetuning strategies. the first one only leverage hard or in-batch negatives, including ance @xcite , seed @xcite , adore @xcite , costa @xcite , prop @xcite , b-prop @xcite , aggretriever @xcite , and co-condener @xcite . the second type leverage sophisticated fine-tuning strategies like knowledge distillation, including rocketqav2 @xcite , ar2 @xcite , ar2+simans @xcite , spladev2 @xcite , colbertv2 @xcite , ernie-search @xcite , simlm @xcite , retromae @xcite . we emphasize two methods for zero-shot evaluations. one is bm25, which is a common sparse retrieval method and a strong baseline in zero-shot settings. the other type are the largescale pre-trained retrievers based on contrastive learning: contriever @xcite and the family of gtr-* @xcite . among them, gtr-xxl is a super large model with 4.8b parameters (over 40× larger than bert base). implementation details. dupmae utilizes a bi-directional transformer network as its encoder, with 12 layers, 768 hidden-dim, and a vocabulary of 30522 tokens (same as bert base). the decoder is an one-layer transformer. the [cls] embedding and ot embedding are reduced to dim-384 by default. as a result, it will preserve the same computation cost of inner-product as the baselines which use dim-768 embeddings. we also explore other configurations of dimensions in our experiments. the masking ratio is set to 0.3 for encoder and 0.5 for decoder. we leverage three commonly used corpora for pre-training: wikipedia, bookcorpus @xcite , and ms marco @xcite . the pre-training and fine-tuning take place on machines with 8× nvidia v100 (32gb) gpus. the models are implemented with pytorch 1.8 and huggingface transformers 4.16. ## conclusion this paper presents dupmae, a new approach for retrieval-oriented pre-training, where the semantic representation capacities can be jointly enhanced for all contextualized embeddings of the language model. it employs two complementary tasks: one reconstructs the original input from the [cls]'s embedding, the other one predicts the bow features based on the ot embeddings. the two tasks are jointly conducted to learn an unified encoder. the two types of embeddings, with reduced dimensions, are aggregated to be a joint semantic representation. the effectiveness of our proposed method is empirically verified, where remarkable performances are achieved on ms marco and beir benchmarks throughout different situations.