R etro MAE : Pre-Training Retrieval-oriented Language Models Via Masked Auto-Encoder (2022)

Abstract: Despite pre-training’s progress in many important NLP tasks, it remains to explore effective pre-training strategies for dense retrieval. In this paper, we propose RetroMAE, a new retrieval oriented pre-training paradigm based on Masked Auto-Encoder (MAE). RetroMAE is highlighted by three critical designs. 1) A novel MAE workflow, where the input sentence is polluted for encoder and decoder with different masks. The sentence embedding is generated from the encoder’s masked input; then, the original sentence is recovered based on the sentence embedding and the decoder’s masked input via masked language modeling. 2) Asymmetric model structure, with a full-scale BERT like transformer as encoder, and a one-layer transformer as decoder. 3) Asymmetric masking ratios, with a moderate ratio for encoder: 15 30%, and an aggressive ratio for decoder: 50 70%. Our framework is simple to realize and empirically competitive: the pre-trained models dramatically improve the SOTA performances on a wide range of dense retrieval benchmarks, like BEIR and MS MARCO. The source code and pre-trained models are made publicly available at https://github.com/staoxiao/RetroMAE so as to inspire more interesting research.

## introduction dense retrieval is important to many web applications. by letting semantically correlated query and document represented as spatially close embeddings, dense retrieval can be efficiently conducted via approximate nearest neighbour search, such as pq @xcite @xcite and hnsw @xcite . recently, large-scale language models have been widely used as the encoding networks for dense retrieval @xcite †. the two researchers make equal contributions to this work and are designated as co-first authors. * . corresponding author. 2020; @xcite . the mainstream models, e.g., bert @xcite , roberta @xcite , t5 @xcite , are usually pre-trained by token-level tasks, like mlm and seq2seq. however, the sentence-level representation capability is not fully developed in these tasks, which restricts their potential for dense retrieval. given the above defect, there have been increasing interests to develop retrieval oriented pretrained models. one popular strategy is to leverage self-contrastive learning @xcite , where the model is trained to discriminate positive samples from data augmentation. however, the self-contrastive learning can be severely limited by the data augmentation's quality; besides, it usually calls for massive amounts of negative samples @xcite . another strategy relies on auto-encoding @xcite @xcite , which is free from the restrictions on data augmentation and negative sampling. the current works are differentiated in how the encoding-decoding workflow is designed, and it remains an open problem to explore more effective auto-encoding framework for retrieval oriented pre-training. we argue that two factors are critical for the autoencoding based pre-training: 1) the reconstruction task must be demanding enough on encoding quality, 2) the pre-training data needs to be fully utilized. we propose retromae (figure 1 ), which optimizes both aspects with the following designs. norwegian forest cat is a breed of dom-estic cat originating in northern europe [m] forest cat is a breed of [m] cat originating in [m] europe [m] [m] cat is [m] [m] of domestic [m] [m] in northern [m] encoder decoder mask (en) mask (de) norwegian forest cat is a breed of domestic cat originating in northern europe sentence embedding figure 1: retromae. the encoder utilizes a full-scale bert, whose input is moderately masked. the decoder is a one-layer transformer, whose input is aggressively masked. the original input is recovered based on the sentence embedding and the decoder's input via mlm. the above designs of retromae are favorable to the pre-training effectiveness thanks to the following reasons. firstly, the auto-encoding is made more demanding on encoding quality. the conventional auto-regression may attend to the prefix during the decoding process; and the conventional mlm only masks a small portion (15%) of the input tokens. by comparison, retromae aggressively masks most of the input for decoding. as such, the reconstruction will be not enough to leverage the decoder's input alone, but heavily depend on the sentence embedding. thus, it will force the encoder to capture in-depth semantics of the input. secondly, it ensures training signals to be fully generated from the input sentence. for conventional mlm-style methods, the training signals may only be generated from 15% of the input tokens. whereas for retromae, the training signals can be derived from the majority of the input. besides, knowing that the decoder only contains onesingle layer, we further propose the enhanced decoding on top of two-stream attention @xcite and position-specific attention mask @xcite . as such, 100% of the tokens can be used for reconstruction, and each token may sample a unique context for its reconstruction. retromae is simple to realize and empirically competitive. we merely use a moderate-amount of data (wikipedia, bookcorpus, ms marco corpus) for pre-training, where a bert base scale encoder is learned. for the zero-shot setting, it produces an average score of 45.2 on beir @xcite ; and for the supervised setting, it may easily reach an mrr@10 of 41.6 on ms marco passage retrieval @xcite following standard knowledge distillation procedures. both values are unprecedented for dense retrievers with the same model size and pre-training conditions. we also carefully evaluate the impact introduced from each of the components, whose results may bring interesting insights to the future research. ## related works dense retrieval is widely applied to web applications, like search engines @xcite , advertising @xcite and recommender systems @xcite . it encodes query and document within the same latent space, where relevant documents to the query can be efficiently retrieved via ann search. the encoding model is critical for the retrieval quality. thanks to the recent development of large-scale language models, e.g., bert @xcite , roberta @xcite , and t5 @xcite , there has been a major leap-forward for dense retrieval's performance @xcite @xcite . the large-scale language models are highly differentiated in terms of pre-training tasks. one common task is the masked language modeling (mlm), as adopted by bert @xcite and roberta @xcite , in which the masked tokens are predicted based on their context. the basic mlm is extended in many ways. for example, tasks like entity masking, phrase masking and span masking @xcite ) may help the pre-trained models to better support the sequence labeling applications, such as entity resolution and question answering. besides, tasks like auto-regression @xcite and seq2seq @xcite are also utilized, where the pre-trained models are enabled to serve nlg related scenarios. however, most of the generic pre-trained models are based on token-level tasks, where the sentence representation capability is not effectively developed @xcite . thus, it may call for a great deal of labeled data @xcite and sophisticated fine-tuning methods @xcite to ensure the pre-trained models' performance for dense retrieval. to mitigate the above problem, recent works propose retrieval oriented pre-trained models. the existing methods can be divided as the ones based on self-contrastive learning (scl) and the ones based on auto-encoding (ae). the scl based methods @xcite @xcite rely on data augmentation, e.g., inverse cloze task (ict), where positive samples are generated for each anchor sentence. then, the language model is learned to discriminate the positive samples from the negative ones via contrastive learning. however, the self-contrastive learning usually calls for huge amounts of negative samples, which is computationally expensive. besides, the pre-training effect can be severely restricted by the quality of data augmentation. the ae based methods are free from these restrictions, where the language models are learned to reconstruct the input sentence based on the sentence embedding. the existing methods utilize various reconstruction tasks, such as mlm @xcite and auto-regression @xcite @xcite , which are highly differentiated in terms of how the original sentence is recovered and how the training loss is formulated. for example, the auto-regression relies on the sentence embedding and prefix for reconstruction; while mlm utilizes the sentence embedding and masked context. the auto-regression derives its training loss from the entire input tokens; however, the conventional mlm only learns from the masked positions, which accounts for 15% of the input tokens. ideally, we expect the decoding operation to be demanding enough, as it will force the encoder to fully capture the semantics about the input so as to ensure the reconstruction quality. besides, we also look forward to high data efficiency, which means the input data can be fully utilized for the pre-training task. ## methodology we develop a novel masked auto-encoder for retrieval oriented pre-training. the model contains two modules: a bert-like encoder φ enc (•) to generate sentence embedding, and a one-layer trans-former based decoder φ dec (•) for sentence reconstruction. the original sentence x is masked as xenc and encoded as the sentence embedding h x . the sentence is masked again (with a different mask) as xdec ; together with h x , the original sentence x is reconstructed. detailed elaborations about retromae are made as follows. ## experimental studies we evaluate the retrieval performance of the sentence embedding generated by retromae's pretrained encoder, where two major issues are explored. 1) retromae's impact on zero-shot and supervised dense retrieval, in comparison with the generic pre-trained models and the retrieval oriented pre-trained models. 2) the impact from the four technical factors in retromae: the enhanced decoding, the decoder size, the decoder's masking ratio, and the encoder's masking ratio. ## conclusion we propose retromae, a novel masked autoencoding framework to pre-train retrieval oriented language models: the input sentence is randomly masked for encoder and decoder, and the sentence embedding is joined with the decoder's masked input to reconstruct the original input. we introduce an asymmetric model structure (full-scale encoder and single-layer decoder) and asymmetric masking ratios (a moderate ratio for encoder and an aggressive one for decoder), which makes the reconstruction sufficiently demanding. we also introduce the enhanced decoding, which makes the full utilization of the pre-training data. our experiments on beir, ms marco, and natural question validate retromae's effectiveness, as significant improvements on both zero-shot and supervised evaluations can be achieved over the existing methods.