Source-linked AI summary
RetroMAE: Pre-Training Retrieval-oriented Language Models Via Masked Auto-Encoder
Shitao Xiao, Zheng Liu, Yingxia Shao, Zhao Cao
TL;DR
Dense retrieval needs pre-training strategies that better develop sentence-level representations than generic token-level objectives. RetroMAE uses masked auto-encoding with asymmetric encoder-decoder structures and masking ratios, and it substantially improves dense retrieval performance across BEIR, MS MARCO, and Natural Question. The method is effective with BERT-base-scale models and moderate pre-training data, but larger models and datasets remain to be explored.
Problem
Generic token-level pre-training does not fully develop sentence-level representations, limiting its potential for dense retrieval.
Method
RetroMAE generates a sentence embedding from one masked input and reconstructs the original sentence from that embedding plus a differently masked decoder input, using asymmetric architecture and masking.
Results
Significant improvements over existing methods are achieved on zero-shot and supervised evaluations across BEIR, MS MARCO, and Natural Question.
Takeaways & Limitations
RetroMAE provides a retrieval-oriented masked auto-encoding framework that is simple to realize and empirically competitive.
Takeaways & Limitations
Experiments use BERT-base-scale transformers and mainly moderate pre-training data, leaving the effects of larger networks and more data unresolved.
Abstract
from arXiv · showhide
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.
1 Introduction
RetroMAE targets the underdeveloped sentence-level representation capability of token-level pre-training for dense retrieval. It combines masked auto-encoding with asymmetric architecture and masking, achieving strong retrieval results under comparable model size and pre-training conditions.
- Token-level pre-training tasks such as MLM and Seq2Seq do not fully develop sentence-level representations needed for dense retrieval.
- RetroMAE masks each input sentence twice, using one masked view to generate a sentence embedding and another to reconstruct the original through MLM.
- The model uses a full-scale BERT encoder and a single-layer transformer decoder for asymmetric representation and reconstruction.
- Encoder masking is moderate at 15∼30%, while decoder masking is aggressive at 50∼70%.
- 45.2 average BEIR score and 41.6 MS MARCO MRR@10 were achieved with moderate pre-training data and a BERT-base-scale encoder.
2 Related works
Related work frames dense retrieval around shared-space encoding and retrieval-oriented pre-training. Existing approaches primarily use self-contrastive learning or auto-encoding, while sentence representations remain insufficiently developed in generic token-level pre-training.
- Dense retrieval encodes queries and documents in a shared latent space for efficient approximate nearest-neighbor search.
- Generic pre-training often develops sentence representations insufficiently, increasing reliance on labeled data and sophisticated fine-tuning for dense retrieval.
- Retrieval-oriented pre-training methods are broadly divided into self-contrastive-learning and auto-encoding approaches.
3 Methodology
RetroMAE uses a BERT-like encoder and lightweight decoder in a masked auto-encoding workflow for retrieval-oriented sentence embeddings. Enhanced decoding diversifies reconstruction contexts while retaining encoder MLM training.
- 3 Methodology: RetroMAE masks the sentence separately for encoding and decoding, generating an embedding from the encoder input and reconstructing the original sentence from the decoder input plus that embedding.The encoder uses a BERT-like architecture, while the decoder reconstructs the sentence with a one-layer transformer.
- 3 Methodology: The encoder applies moderate masking of 15∼30%, preserving most input information while producing the sentence embedding from the polluted sentence.A 12-layer, 768-dimensional BERT-like encoder uses the final [CLS] hidden state as the sentence embedding.
- 3 Methodology: The decoder applies 50∼70% masking and combines the masked input with the sentence embedding to reconstruct the original sentence using cross-entropy over masked tokens.The one-layer decoder and aggressive masking make reconstruction challenging and encourage high-quality sentence embeddings.
- 3 Methodology: Enhanced decoding uses query and context streams with a position-specific attention mask, allowing each token to be reconstructed from the context visible in its mask row.The output of attention, together with the query stream through residual connection, reconstructs the original input.
- 3 Methodology: The enhanced mask exposes sampled tokens and the first position while excluding each token’s diagonal position, so tokens cannot attend to themselves during reconstruction.The diagonal is filled with −∞ and visible-context positions with 0 in the workflow illustration.
- 3 Methodology: RetroMAE adds the encoder MLM loss to the decoder loss, avoids sophisticated data augmentation and negative sampling, and keeps computation similar to conventional BERT/RoBERTa pre-training.The decoder’s simplicity helps maintain comparable computation cost.
1 begin
The training procedure updates the model using the combined encoder and decoder losses. Its design avoids sophisticated data augmentation and negative sampling while keeping computation similar to conventional BERT/RoBERTa pre-training.
- 1 begin: The model is updated with respect to the combined encoder and decoder losses, Lenc + Ldec.
- 1 begin: RetroMAE requires no sophisticated data augmentation or negative sampling and maintains computation similar to conventional BERT/RoBERTa-style pre-training.This follows from the simplicity of the decoder.
4 Experimental Studies
RetroMAE is evaluated for zero-shot transfer, supervised dense retrieval, and the effects of four technical design choices. It achieves strong results across these evaluations, while ablations support enhanced decoding, a one-layer decoder, and asymmetric masking ratios.
- Main Results: On MS MARCO with ANCE fine-tuning, RetroMAE outperforms coCondenser by +1.1% on MRR@10 using the same model size and pre-training data.
- Main Results: Knowledge distillation results show RetroMAE surpassing AR2 by +2.1%, RocketQAv2 by +2.8%, ERNIE-search by +1.5%, and ColBERTv2 by +1.9%.
- Ablation Studies: Encoder masking that is too aggressive severely harms retrieval because it discards useful sentence information needed for high-quality embeddings.
- Ablation Studies: Ablations find that enhanced decoding improves performance, a one-layer decoder is best, and aggressive decoder plus moderately increased encoder masking improve retrieval quality.
5 Conclusion
RetroMAE is a masked auto-encoding framework for retrieval-oriented language-model pre-training that combines asymmetric architecture and masking with enhanced decoding. Experiments validate significant improvements over existing methods in both zero-shot and supervised dense-retrieval evaluations.
- RetroMAE masks the input separately for encoder and decoder, then reconstructs the original sentence from the sentence embedding and decoder input.
- The framework uses a full-scale encoder, a single-layer decoder, and moderate encoder masking with aggressive decoder masking.
- Enhanced decoding is introduced to make fuller use of pre-training data.
- Significant improvements are achieved over existing methods on BEIR, MS MARCO, and Natural Question in zero-shot and supervised evaluations.
6 Limitations
The empirical studies use BERT base-scale transformers and only a moderate amount of pre-training data. The effects of larger networks and more data therefore remain to be explored.
- The experiments are limited to BERT base-scale transformers and a moderate amount of pre-training data.
- The impact of enlarged networks and increased pre-training data remains unexplored.