[论文解读] RetroMAE: Pre-Training Retrieval-oriented Language Models Via Masked Auto-Encoder

Shitao Xiao, Zheng Liu+2 more|arXiv (Cornell University)|May 24, 2022

Topic Modeling被引用 7

一句话总结

RetroMAE 提出了一种用于预训练检索导向语言模型的新型掩码自编码框架，采用非对称设计，包含一个全编码器 BERT 和单层解码器，并对解码器使用激进的掩码（50–70%），对编码器使用适中掩码（15–30%）。该设置迫使编码器学习深层语义表征，在 BEIR 和 MS MARCO 基准测试中以极少的数据和计算量达到最先进性能。

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.

研究动机与目标

• 为解决现有预训练方法在构建用于密集检索的强句子级表征方面的局限性。
• 克服基于自编码的检索预训练中数据效率和表征质量的问题。
• 设计一种框架，以最大化输入句子的训练信号利用，同时增加编码器的重建难度。
• 探索非对称模型结构和掩码比率在检索导向预训练中的有效性。

提出的方法

• RetroMAE 采用一种新颖的掩码自编码流程，其中输入句子在编码器和解码器上以不同方式被掩码。
• 编码器为全规模 BERT，从适度掩码的输入（15–30% 掩码）生成句子嵌入。
• 解码器为单层 Transformer，利用句子嵌入和高度掩码的输入（50–70% 掩码）重建原始句子。
• 通过双流注意力和位置特定注意力掩码引入增强解码，以确保所有标记均被用于重建并获得独特上下文。
• 采用非对称掩码比率，以增加解码器的重建难度，同时保持编码器输入质量。
• 模型在 Wikipedia、BookCorpus 和 MS MARCO 上进行预训练，重点在于最大化信号利用和表征质量。

实验结果

研究问题

• RQ1 具有非对称掩码比率的掩码自编码框架是否能提升密集检索的句子表征质量？
• RQ2 使用单层解码器并结合增强解码如何影响数据效率和检索性能？
• RQ3 在检索导向的预训练设置中，编码器和解码器的最佳掩码比率是什么？
• RQ4 非对称模型架构（全编码器，轻量解码器）是否优于对称设计？
• RQ5 增强解码在多大程度上提升了检索预训练中训练信号的多样性与模型泛化能力？

主要发现

• RetroMAE 在 BEIR 基准测试中实现 45.2 的零样本平均得分，显著优于现有方法。
• 在 DPR 上进行有监督微调时，RetroMAE 在多个密集检索基准上达到最先进性能。
• 结合双流注意力和位置特定掩码的增强解码提升了数据效率，并带来显著的性能提升。
• 单层解码器优于更深的解码器，因为更大的解码器需要禁用增强解码并导致性能下降。
• 编码器掩码比率提高（最高至 30%）和解码器激进掩码（50–70%）可提升检索质量，最优解码器掩码比为 0.5。
• 消融实验确认，非对称掩码和增强解码对于最大化训练信号利用和模型有效性至关重要。