PaperlayerView exact revision on arXiv

Source-linked AI summary

ConTextual Masked Auto-Encoder for Dense Passage Retrieval

Xing Wu, Guangyuan Ma, Meng Lin, Zijia Lin, Zhongyuan Wang, Songlin Hu

arXiv:2208.07670v3cs.CLcs.AI

TL;DR

Dense passage retrieval relies on strong text representations, while prior reconstruction pre-training focused on internal text modeling and did not account for contextual span relationships. CoT-MAE uses self- and context-supervised masked auto-encoding over neighboring spans, and reports gains across retrieval benchmarks, including over coCondenser on MS-MARCO MRR@10.

  • Problem

    Dense retrieval depends on high-quality text representations, while prior reconstruction pre-training focused on internal text modeling without considering contextual information between spans.

    Source: Introduction · p. 1

  • Method

    CoT-MAE pairs neighboring spans and uses self-supervised reconstruction plus context-supervised reconstruction conditioned on the neighboring span's encoded embedding.

    Source: Introduction · p. 2

  • Results

    CoT-MAE reports gains over competing baselines across MS-MARCO, TREC DL, and Natural Questions, exceeding coCondenser by 1.2% on MS-MARCO MRR@10.

    Source: Introduction · p. 2, Main Results · p. 5

  • Takeaways & Limitations

    The results support generative pre-training that models both span semantics and neighboring-span correlations as an approach for dense passage retrieval.

    Source: Conclusions and Future work · p. 7

Abstract

from arXiv · show

Dense passage retrieval aims to retrieve the relevant passages of a query from a large corpus based on dense representations (i.e., vectors) of the query and the passages. Recent studies have explored improving pre-trained language models to boost dense retrieval performance. This paper proposes CoT-MAE (ConTextual Masked Auto-Encoder), a simple yet effective generative pre-training method for dense passage retrieval. CoT-MAE employs an asymmetric encoder-decoder architecture that learns to compress the sentence semantics into a dense vector through self-supervised and context-supervised masked auto-encoding. Precisely, self-supervised masked auto-encoding learns to model the semantics of the tokens inside a text span, and context-supervised masked auto-encoding learns to model the semantical correlation between the text spans. We conduct experiments on large-scale passage retrieval benchmarks and show considerable improvements over strong baselines, demonstrating the high efficiency of CoT-MAE. Our code is available at https://github.com/caskcsg/ir/tree/main/cotmae.

Introduction

Dense retrieval depends on strong text representations, but prior self-supervised reconstruction work focused on within-text modeling, while contrastive span prediction modeled relationships between spans. CoT-MAE uses paired neighboring spans and self- and context-supervised masked auto-encoding to train representations; experiments report gains over competing retrievers.

Introduction

  • Prior reconstruction methods modeled text internally, while contrastive span-prediction methods modeled semantic relationships between spans; CoT-MAE asks whether contextual span-level generative pre-training can help dense retrieval.The cited reconstruction work includes (Gao and Callan 2021a; Liu and Shao 2022), and the contrastive work includes (Chang et al. 2020; Gao and Callan 2021b; Ma et al. 2022).Source: Introduction · p. 1
  • CoT-MAE forms pairs from neighboring document spans and jointly models within-span token semantics and semantic correlations between spans.Each span serves as the other's context during context-supervised reconstruction.Source: Introduction · pp. 1–2
  • Self-supervised masked auto-encoding reconstructs a span from its unmasked tokens, while context-supervised auto-encoding also conditions reconstruction on the neighboring span's encoded embedding.The decoder is used for pre-training but removed during dense-retrieval fine-tuning.Source: Introduction · p. 2
  • A deep encoder and shallow decoder are jointly trained with asymmetric masking, using rates up to 30% and 45%, respectively, to strengthen encoder representations.The larger decoder-side mask rate makes context-supervised reconstruction harder.Source: Introduction · p. 2
  • CoT-MAE achieves considerable gains over competing baseline retrievers on MS-MARCO, TREC DL, and Natural Questions.The paper reports extensive ablations to examine its design.Source: Introduction · p. 2

Related Works

Related work describes two pre-training approaches for dense retrieval: auxiliary auto-encoding and contrastive span prediction. It also reviews fine-tuning methods, while this paper evaluates pre-training using the Tevatron fine-tuning pipeline.

Pre-training tasks tailored for dense retrieval

Fine-tuning

  • Fine-tuning work includes hard-negative mining, late interaction, distillation, query clustering, data augmentation, and joint retriever-reranker optimization.To isolate the contribution of pre-training, the authors reuse the Tevatron fine-tuning pipeline (Gao et al. 2022b).Source: Related Works · p. 3

Approach

CoT-MAE uses masked auto-encoding to model both token semantics within text spans and semantic correlations between spans. Its approach pairs neighboring spans, trains an asymmetric encoder-decoder, then uses the retained encoder for dense retrieval.

Preliminary: Masked Auto-Encoding

  • Masked auto-encoding prepends [CLS], masks a selected proportion of tokens, and trains the encoder to recover those tokens from the unmasked text.The paper uses a mask rate greater than BERT's 15% setting.Source: Preliminary: Masked Auto-Encoding · p. 3

CoT-MAE: ConTextual Masked Auto-Encoder

  • CoT-MAE pairs neighboring text spans and jointly models token semantics within a span and semantic correlations between spans.The paired spans are sampled from documents; Near, Olap, and Rand respectively form adjacent non-overlapping, adjacent partially overlapping, and random non-overlapping pairs.Source: CoT-MAE: ConTextual Masked Auto-Encoder · pp. 3–4
  • A strong deep encoder and weak shallow decoder use asymmetric masking to reconstruct masked tokens on their respective sides.The encoder reconstructs masked tokens from the unmasked tokens in TA; the decoder reconstructs masked tokens in TB using its unmasked tokens and the neighboring span's context embedding.Source: CoT-MAE: ConTextual Masked Auto-Encoder · p. 4
  • The context-supervised decoder loss and encoder-side self-supervised loss are combined, with the paired spans also modeled in the dual direction.The context embedding for TB is the last-layer [CLS] state of TA; the described loss combines encoder and decoder losses.Source: CoT-MAE: ConTextual Masked Auto-Encoder · p. 4

Fine-tuning on Dense Passage Retrieval

  • The encoder is initialized for query and passage retrieval using its last-layer [CLS] embedding, and the query-passage similarity is an inner product.The decoder is discarded after pre-training; query and passage encoders are fine-tuned with a contrastive loss, using BM25 negatives and first-stage-retriever-mined hard negatives in the stated evaluation pipeline.Source: Fine-tuning on Dense Passage Retrieval · p. 4

Experiments

CoT-MAE is evaluated on passage-ranking and open-domain QA benchmarks, with its main results showing gains over competing baselines across MS-MARCO, NQ, and TREC DL.

Pre-training

  • Pre-training uses 3.2M MS-MARCO documents, a uniform mixture of three span-sampling strategies, and up to 1,200k steps; the decoder is discarded before fine-tuning.The encoder starts from 12-layer BERT-base, while the decoder is initialized from scratch; training uses 8 A100 GPUs for four days, and hyperparameters are not tuned due to the compute budget.Source: Pre-training · p. 4

Fine-tuning

  • Fine-tuning evaluates CoT-MAE on MS-MARCO, NQ, and TREC DL using Tevatron with seed 42, without distillation or multi-vector methods.Reported metrics are MRR@10, R@50, and R@1000 for MS-MARCO; R@5, R@20, and R@100 for NQ; and nDCG@10 for TREC DL.Source: Fine-tuning · pp. 4–5
  • Baselines include sparse retrievers such as BM25 and dense retrievers such as ANCE, SEED, TAS-B, and RetroMAE.Results are drawn mainly from prior published work and include reproduced results with public checkpoints, as indicated in Table 1.Source: Fine-tuning · p. 5, Table 1 · p. 5

Main Results

  • CoT-MAE exceeds coCondenser by 1.2% on MS-MARCO MRR@10 and reaches 70 on TREC DL nDCG@10; on NQ, it is comparable to coCondenser.CoT-MAE also outperforms the evaluated baselines on TREC DL; on NQ, it beats the remaining methods, while coCondenser is slightly better on R@5 and CoT-MAE on R@100.Source: Main Results · p. 5, Table 1 · p. 5

Analysis

The analysis compares CoT-MAE with distilled retrievers and probes design choices, including masking, span sampling, decoder depth, and decoder initialization. Unless otherwise stated, analyses use MS-MARCO passage ranking with 800k pre-training steps rather than the main experiment’s 1200k steps.

Comparison with Distilled Retrievers

  • After distillation, CoT-MAE exceeds all previously state-of-the-art distilled retrievers on MRR@10, R@50, and R@1k.Without distillation, it slightly outperforms all distilled retrievers on R@1k, trails only AR2 on R@50, and surpasses RocketQAv2 on MRR@10.Source: Table 2 · p. 6, Comparison with Distilled Retrievers · p. 6

Impact of Mask Rate

  • The best mask-rate setting uses 30% encoder masking and 45% decoder masking; performance is robust across a wide range of rates.At a 45% encoder mask rate, performance decreases slightly, which the authors attribute to insufficient encoder context.Source: Table 3 · p. 6, Impact of Mask Rate · p. 6

Impact of Sampling Strategies

  • Mixing Near, Olap, and Rand sampling strategies gives the best performance, while using a single strategy causes a larger drop than mixing two.The authors interpret the trend as showing that diversity in data composition can benefit pre-training.Source: Table 4 · p. 6, Impact of Sampling Strategies · p. 6

Impact of Decoder Layer Number

  • A two-layer decoder performs best; weaker or stronger decoders reduce performance.A one-layer decoder may not fuse context and unmasked-token embeddings fully, while a deeper decoder relies less on context and constrains encoder training less.Source: Table 5 · p. 6, Impact of Decoder Layer Number · pp. 6–7

Qualitative Analysis

  • In qualitative examples, CoT-MAE retrieves more semantically relevant rank-1 passages than BERT or coCondenser, whose results can overlap with query tokens without being highly relevant.The authors connect this difference to joint masked-autoencoding of token and span context.Source: Table 6 · p. 7, Qualitative Analysis · p. 7

Impact of Decoder Weight Initialization

Conclusions and Future work

CoT-MAE is a generative pre-training method for dense retrieval that models token semantics within spans and semantic correlations between neighboring spans. The paper reports considerable gains over competing baseline retrievers and identifies multilingual and multimodal applications as future directions.

Loading 2208.07670v3…