RetroMAE-2: Duplex Masked Auto-Encoder For Pre-Training Retrieval-Oriented Language Models

TL;DR

RetroMAE-2 introduces a duplex auto-encoding pre-training method that enhances semantic representations by jointly leveraging all token embeddings, significantly improving retrieval performance on benchmarks like MS MARCO and BEIR.

Contribution

It proposes Duplex Masked Auto-Encoder (DupMAE), a novel pre-training approach that jointly trains on two tasks to utilize all contextualized embeddings for retrieval tasks.

Findings

Improves semantic representation quality for retrieval models.

Achieves superior performance on MS MARCO and BEIR benchmarks.

Enhances transferability of pre-trained language models.

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.