信息检索（113）：RetroMAE: Pre-Training Retrieval-oriented Language Models Via Masked Auto-Encoder

---

发布时间（2022）

---

通过掩蔽自动编码器预训练面向检索的语言模型

摘要

尽管预训练在许多重要的 NLP 任务中取得了进展，但仍需探索用于密集检索的有效预训练策略。 在本文中，我们提出了 RetroMAE，一种基于掩码自动编码器 (MAE) 的面向检索的新型预训练范式。RetroMAE 有三个关键设计。1）一种新颖的 MAE 工作流程，其中输入句子被不同的掩码污染用于编码器和解码器。句子嵌入由编码器的掩码输入生成；然后，通过掩码语言建模，基于句子嵌入和解码器的掩码输入恢复原始句子。2）非对称模型结构，以全尺寸 BERT 类 Transformer 作为编码器，以单层 Transformer 作为解码器。 3）非对称掩码率，编码器的比率适中：15∼30%，解码器的比率激进：50∼70%。我们的框架易于实现且具有经验竞争力：预训练模型显著提高了各种密集检索基准（如 BEIR 和 MS MARCO）上的 SOTA 性能。源代码和预训练模型已在 https://github.com/staoxiao/RetroMAE 上公开发布，以激发更多有趣的研究。

1 引言

密集检索对许多 Web 应用来说都很重要。通过将语义相关的查询和文档表示为空间上接近的嵌入，可以通过近似最近邻搜索有效地进行密集检索，例如 PQ (Jegou et al, 2010; Xiao et al, 2021, 2022a) 和 HNSW (Malkov and Yashunin, 2018)。 最近，大规模语言模型已被广泛用作密集检索的编码网络 (Karpukhin et al, 2020; Xiong et al, 2020; Luan et al, 2021)。主流模型，例如 BERT (Devlin et al, 2019)、RoBERTa (Liu et al, 2019)、T5 (Raffel et al, 2019)，通常由 token 级任务（如 MLM 和 Seq2Seq）进行预训练。然而，句子级的表征能力在这些任务中还没有得到充分开发，这限制了它们进行密集检索的潜力

鉴于上述缺陷，人们对开发面向检索的预训练模型的兴趣日益浓厚。一种流行的策略是利用自对比学习 (Chang et al, 2020; Guu et al, 2020)，其中训练模型以区分正样本和数据增强。然而，自对比学习可能受到数据增强质量的严重限制；此外，它通常需要大量的负样本 (He et al, 2020a; Chen et al, 2020)。 另一种策略依赖于自动编码 (Gao and Callan, 2021; Lu et al, 2021; Wang et al, 2021)，它不受数据增强和负采样的限制。当前的研究在编码-解码工作流程的设计上有所不同，探索更有效的面向检索的预训练自动编码框架仍然是一个悬而未决的问题。

我们认为基于自动编码的预训练有两个关键因素：1）重建任务必须对编码质量有足够高的要求；2）预训练数据需要得到充分利用。我们提出了 RetroMAE（图 1），它通过以下设计优化了这两个方面。

• 一种新颖的 MAE 工作流程。预训练遵循一种新颖的掩码自动编码工作流程。输入句子被两个不同的掩码污染两次。一个掩码输入由编码器使用，其中生成句子嵌入。另一个由解码器使用：与句子嵌入结合，通过掩码语言建模 (MLM) 恢复原始句子。

• 非对称结构。RetroMAE 采用非对称模型结构。编码器是全尺寸 BERT，能够为输入句子生成判别性嵌入。相比之下，解码器遵循极其简化的结构，即单层 Transformer，通过学习来重构输入句子。

• 不对称掩蔽率。编码器的输入被掩蔽的比例适中：15∼30%，略高于 MLM 中的传统值。然而，解码器的输入被掩蔽的比例则要高得多：50∼70%。

RetroMAE 的上述设计有利于提高预训练的效果，原因如下。首先，自动编码对编码质量的要求更高。传统的自回归可能会在解码过程中关注前缀；而传统的 MLM 只会掩盖一小部分（15%）输入标记。相比之下，RetroMAE 会在解码时积极掩盖大部分输入。因此，仅靠重构不足以利用解码器的输入，而是严重依赖于句子嵌入。因此，它将迫使编码器捕获输入的深层语义。其次，它确保训练信号完全从输入句子中生成。对于传统的 MLM 式方法，训练信号可能仅从 15% 的输入标记中生成。而对于 RetroMAE，训练信号可以从大部分输入中得出。此外，考虑到解码器仅包含一个层，我们进一步提出了在双流注意力（Yang et al，2019）和位置特定注意力掩码（Dong et al，2019）基础上的增强解码。这样，100% 的 token 都可以用于重建，并且每个 token 都可以为其重建采样一个唯一的上下文。

RetroMAE 实现简单，经验上具有竞争力。我们仅使用中等量的数据（Wikipedia、BookCorpus、MS MARCO 语料库）进行预训练，其中学习了 BERT 基础尺度编码器。对于零样本设置，它在 BEIR（Thakur et al，2021）上产生 45.2 的平均分数；对于监督设置，按照标准知识蒸馏程序，它可以在 MS MARCO 段落检索（Nguyen et al，2016）上轻松达到 41.6 的 MRR@10。对于具有相同模型大小和预训练条件的密集检索器，这两个值都是前所未有的。 我们还仔细评估了每个组件引入的影响，其结果可能会为未来的研究带来有趣的见解。

2 相关工作

密集检索广泛应用于网络应用，如搜索引擎（Karpukhin 等人，2020 年）、广告（Lu 等人，2020 年；Zhang 等人，2022 年）和推荐系统（Xiao 等人，2022b 年）。它在同一个潜在空间内对查询和文档进行编码，通过 ANN 搜索可以有效地检索与查询相关的文档。编码模型对于检索质量至关重要。 得益于最近大规模语言模型的发展，例如 BERT（Devlin 等人，2019 年）、RoBERTa（Liu 等人，2019 年）和 T5（Raffel 等人，2019 年），密集检索的性能有了重大飞跃（Karpukhin 等人，2020 年；Luan 等人，2021 年；Lin 等人，2021 年）。

大规模语言模型在预训练任务方面差异很大。一个常见的任务是掩码语言模型（MLM），如 BERT（Devlin et al, 2019）和 RoBERTa（Liu et al, 2019）所采用的，其中根据上下文预测掩码标记。基本的 MLM 以多种方式扩展。例如，实体掩码、短语掩码和跨度掩码（Sun et al, 2019; Joshi et al, 2020）等任务可以帮助预训练模型更好地支持序列标记应用，例如实体解析和问答。此外，自回归（Radford et al, 2018; Yang et al, 2019）和 Seq2Seq（Raffel et al, 2019; Lewis et al, 2019）等任务也被使用，其中预训练模型可用于 NLG 相关场景。然而，大多数通用预训练模型都是基于 token 级任务的，其中句子表征能力尚未得到有效开发（Chang et al, 2020）。因此，可能需要大量标记数据（Nguyen et al, 2016; Kwiatkowski et al, 2019）和复杂的微调方法（Xiong et al, 2020; Qu et al, 2020）来确保预训练模型在密集检索中的性能。

为了缓解上述问题，最近的研究提出了面向检索的预训练模型。现有的方法可以分为基于自对比学习（SCL）的方法和基于自动编码（AE）的方法。基于 SCL 的方法（Chang 等人，2020；Guu 等人，2020；Xu 等人，2022）依赖于数据增强，例如逆完形填空任务（ICT），其中为每个锚句子生成正样本。然后，通过对比学习学习语言模型以区分正样本和负样本。然而，自对比学习通常需要大量的负样本，这在计算上是昂贵的。此外，数据增强的质量会严重限制预训练效果。基于 AE 的方法不受这些限制，其中语言模型被学习以基于句子嵌入重建输入句子。 现有的方法利用各种重建任务，例如 MLM（Gao and Callan，2021）和自回归（Lu et al，2021；Wang et al，2021；Li et al，2020），它们在如何恢复原始句子以及如何制定训练损失方面存在高度差异。例如，自回归依赖于句子嵌入和前缀进行重建；而 MLM 利用句子嵌入和屏蔽上下文。自回归从整个输入标记中得出其训练损失；然而，传统的 MLM 只从屏蔽位置学习，这占输入标记的 15%。理想情况下，我们希望解码操作足够苛刻，因为它将迫使编码器完全捕获有关输入的语义，以确保重建质量。此外，我们还期待高数据效率，这意味着输入数据可以充分利用于预训练任务。

3 方法

我们开发了一种新型的掩码自动编码器，用于面向检索的预训练。该模型包含两个模块：一个类似 BERT 的编码器 Φenc(·) 用于生成句子嵌入，以及一个基于单层 Transformer 的解码器 Φdec(·) 用于句子重构。原始句子 X 被掩码为 ˜Xenc 并编码为句子嵌入 hX˜ 。 该句子再次被掩码（使用不同的掩码）为 ˜Xdec；与 hX˜ 一起，原始句子 X 被重构。关于 RetroMAE 的详细阐述如下。

3.1 编码

输入句子 X 在编码阶段被污染为 ˜Xenc，其中一小部分标记被特殊标记 [M] 随机替换（图 2. A）。我们应用适度的掩蔽率（15∼30%），这意味着输入的大部分信息将被保留。然后，编码器 Φ enc(·) 用于将污染的输入转换为句子嵌入 hX˜：

我们采用了类似 BERT 的编码器，该编码器有 12 层和 768 个隐藏维度，有助于捕捉句子的深度语义。按照惯例，我们选择 [CLS] 标记的最终隐藏状态作为句子嵌入。

3.2 解码

输入句子 X 在解码阶段被污染为 ˜Xdec（图 2. B）。掩蔽率比编码器使用的掩蔽率更高，其中 50∼70% 的输入标记将被掩蔽。掩蔽的输入与句子嵌入相结合，解码器在此基础上重建原始句子。具体而言，句子嵌入和掩蔽的输入组合成以下序列：

在上式中，exi 表示 xi 的嵌入，其中添加了一个额外的位置嵌入 pi。最后，通过优化以下目标，解码器 Φdec 学习重建原始句子 X：

其中 CE 是交叉熵损失。如前所述，我们使用基于单层 Transformer 的解码器。鉴于输入被大量掩蔽且网络被极度简化，解码变得具有挑战性，这迫使生成高质量的句子嵌入，以便能够以良好的保真度恢复原始输入。

3.3 增强解码

解码过程的一个限制是训练信号（即交叉熵损失）只能从被掩蔽的 token 中得出。此外，每个被掩蔽的 token 总是基于相同的上下文（即 HX˜dec）进行重构。我们认为，只要 1）可以从输入句中得出更多的训练信号，2）可以基于多样化的上下文执行重构任务，就可以进一步增强预训练效果。为此，我们提出了受双流自注意力（Yang et al, 2019）和位置特定注意力掩码（Dong et al, 2019）启发的增强解码。具体而言，我们为解码操作生成两个输入流：H1（查询）和 H2（上下文）（图 2.C）：

其中 hX˜ 是句子嵌入，exi 是 token 嵌入（此处没有 token 被掩码），pi 是位置嵌入。我们引入位置特定的注意力掩码 M ∈ R L×L，其中自注意力机制如下：

输出 A 与 H1（由于残差连接）一起用于重建原始输入（其他操作，如 layer-norm 和 FFN，我们在讨论中省略了）。最后，将优化以下目标：

已知解码器仅由一个单独的 Transformer 层组成，每个 token xi 基于矩阵 M 第 i 行可见的上下文进行重构。在这里，应用以下规则来生成特定位置的注意力掩码矩阵 M：

重建 xi 时，采样的标记 s(X6=i) 和第 1 个位置（第 1 行除外）将可见。
对角线元素，即第 i 行的 xi，将始终被排除，这意味着它们将始终被屏蔽；因此，每个标记在重建过程中都无法关注自身。

我们将增强解码的预训练工作流程总结为算法 1。注意，BERT 中原始的掩蔽语言建模任务保留在编码器中。相应的损失（表示为 Lenc）与解码器的损失相加，形成最终的训练损失。我们的预训练工作流程需要强调以下特点。首先，由于解码器的掩蔽率很高且网络极其简化，因此重建任务非常艰巨。 其次，我们可以从无监督语料库中获取丰富的预训练信号，因为每个输入句子中的所有标记都可以用于重建。最后，预训练很容易实现：1）不需要复杂的数据增强和负采样，2）由于解码器简单，计算成本与传统的 BERT/RoBERTa 风格的预训练保持相似。

4 实验研究

我们评估了 RetroMAE 预训练编码器生成的句子嵌入的检索性能，其中探讨了两个主要问题。1）与通用预训练模型和面向检索的预训练模型相比，RetroMAE 对零样本和监督密集检索的影响。2）RetroMAE 中的四个技术因素的影响：增强解码、解码器大小、解码器的掩蔽率和编码器的掩蔽率。

4.1 实验设置

以下数据集用于 RetroMAE 的预训练和评估
• 预训练。我们重复使用与 BERT（Devlin 等人，2019 年）相同的预训练语料库：英文维基百科和 BookCorpus。以前的面向检索的预训练工作也经常使用相应的数据集，例如 SEED（Lu 等人，2021 年）和 Condenser（Gao 和 Callan，2021 年）。继 coCondenser（Gao 和 Callan，2022 年）之后，我们还使用 MS MARCO 语料库来分析域内预训练对密集检索的影响（我们发现这对 MS MARCO 上的性能至关重要，但对其他数据集上的性能则不必要）。

• 评估。我们使用两个数据集来评估监督后的检索性能。1) MS MARCO (Nguyen et al, 2016)，其中包含来自 Bing 搜索的查询。我们使用其段落检索任务，其中包含 502,939 个训练查询和 6,980 个评估查询 (Dev)。 需要从 880 万个候选段落中检索答案。2) Natural Questions (Kwiatkowski et al, 2019)，由来自 Google 的查询组成。有 79,168 个训练查询、8,757 个开发查询和 3,610 个测试查询。答案是从 21,015,324 个维基百科段落中检索出来的。我们在 BEIR 基准 (Thakur et al, 2021) 上评估零样本检索性能。它使用 MS MARCO 查询对预训练模型进行微调，并在其他 18 个数据集上评估零样本可迁移性。评估数据涵盖不同领域的密集检索任务，例如问答、事实核查、生物医学检索、新闻检索和社交媒体检索等。 我们在实验研究中考虑了三种类型的基线方法1

• 通用模型。我们在实验中使用了以下通用预训练语言模型。1）BERT（Devlin 等人，2019），这是实践中最流行的预训练语言模型。它也经常被微调为密集检索器的编码网络（Karpukhin 等人，2020；Xiong 等人，2020）。2）RoBERTa（Liu 等人，2019），它是 BERT 的增强复制，具有大量扩大的训练数据和改进的训练设置。3）DeBERTa（He 等人，2020b），它通过解缠注意力机制和增强的掩码解码器进一步改进了 BERT 和 RoBERTa；它是 NLU 基准（例如 GLUE 和 SuperGLUE）上最强大的预训练模型之一。

• 面向检索的模型。我们在实验中考虑了两种面向检索的预训练模型。一种是基于自对比学习的，包括以下方法。1）SimCSE（Gao et al，2021），其中学习语言模型以区分通过 drop-out 增强的锚句子的不同视图。2）LaPraDoR（Xu et al，2022），传统 ICT 预训练的增强（Guu et al，2020；Chang et al，2020）；它提出迭代训练查询和文档编码器，以便可以扩大负样本的规模。3）DiffCSE（Chuang et al，2022），它通过联合利用自对比学习和条件差异预测来增强 SimCSE。另一种是基于自动编码的，包括以下方法。 4）Condenser（Gao and Callan，2021），其中句子嵌入与来自编码器的中间隐藏状态相结合，用于MLM。5）SEED（Lu et al，2021），其中句子嵌入用于通过自回归恢复原始输入。

实现细节。RetroMAE 使用双向 Transformer 作为编码器，具有 12 层、768 个隐藏层和 30522 个 token 词汇表（与 BERT 基数相同）。解码器是单层 Transformer。编码器的默认掩码率为 0.3，解码器的默认掩码率为 0.5。该模型训练了 8 个 epoch，使用 AdamW 优化器，批处理大小为 32（每个设备），学习率为 1e-4。训练在一台配备 8× Nvidia A100（40GB）GPU 的机器上进行。这些模型是使用 PyTorch 1.8 和 HuggingFace Transformer 4.16 实现的。 我们采用 BEIR 的官方脚本 2 来准备模型进行零样本评估。对于监督评估，我们使用 DPR（Karpukhin 等人，2020 年）和 ANCE（Xiong 等人，2020 年）来微调预训练模型。我们还使用标准知识蒸馏来评估模型在 MS MARCO 上的性能。具体来说，我们利用 ANCE 微调双编码器返回的硬负样本训练一个 BERT 基础尺度交叉编码器；然后，我们通过使用交叉编码器最小化 KL 散度来进一步微调双编码器。

4.2 主要结果

我们分析了表 1 中的零样本性能，其中 RetroMAE 取得了显著的优势：它在大多数数据集上产生了最好的经验结果，并且总体平均值比最强的基线高出 4.5%。零样本性能的飞跃远远大于基线模型所做的任何改进。 这种改进确实令人兴奋，因为它不是来自模型规模的增加或预训练数据的丰富，而纯粹来自预训练算法的升级。我们进一步在表 2 和表 3 中展示了监督评估，其中预训练模型使用 DPR 和 ANCE 进行了微调。根据基线是通用预训练模型还是面向检索的模型，将基线分为两组。可以观察到，RetroMAE 相对于基线保持了显著的优势。通过 DPR 微调，它在两个数据集上的表现比最强基线高出 +1.96%（MRR@10）和 +1.48%（Recall@10）；通过 ANCE 微调，相应的优势变为 +1.42%（MRR@10）和 +1.41%（Recall@10）。

至于 RetroMAE 在 MS MARCO 语料上的预训练模型。通过 ANCE 微调（表 4），我们的方法在 MRR@10（具有相同的模型大小和相同的预训练数据）上比同类方法 coCondenser（Gao and Callan，2022）高出 +1.1%。而在知识蒸馏场景下（表 5），RetroMAE 超越了近期提出的一系列强基线，包括具有高度复杂蒸馏方法的模型：AR2（Zhang et al，2021）高出 +2.1%，RocketQAv2（Ren et al，2021）高出 +2.8%，ERNIE-search（Lu et al，2022）高出 +1.5%，以及后期交互模型 ColBERTv2（Santhanam et al，2021）高出 +1.9%。

总结一下，以上结果从两个方面验证了 RetroMAE 的有效性。1）它显著提高了预训练模型的可迁移性，有助于在域外数据集上实现出色的零样本性能。2）它为密集检索器提供了强大的初始化；在使用域内数据进行微调后，它在相应的场景中产生了高质量的监督检索性能。除了主要结果之外，我们还可能有以下有趣的观察结果。

首先，通用领域的高级预训练方法不一定能提高密集检索性能。特别是，RoBERTa 和 DeBERTa 都是 BERT 的重大改进；然而，它们都无法像在通用 NLU 基准上那样取得比 BERT 更好的性能。这一观察结果进一步支持了开发面向检索的预训练模型的动机。 其次，自动编码风格的预训练（由 SEED、Condenser 和 RetroMAE 采用）被经验证明更有利于密集检索，因为它在零样本评估和监督评估中都优于通用和基于自对比学习的预训练模型。

第三，基于自对比学习的预训练模型在微调时与通用模型相比几乎没有什么改进，因为它们在监督和零样本评估中的表现都很接近。事实上，考虑到最近在密集检索（Gao and Callan，2021）（BERT 对抗 ICT）和图像处理（ElNouby et al，2021）（BEiT（Bao et al，2021）对抗 MoCo/DINO）方面的研究报告的类似结果，这一观察结果没有任何监督。也就是说，来自自对比学习的预训练模型往往具有相对有限的微调潜力。因此，我们认为面向检索的预训练算法应该根据具体场景中的微调情况进行适当设计。

4.3 消融研究

我们在表 6 中对 RetroMAE 及其监督性能（DPR 微调）进行了分析，其中分析了以下因素：1）解码方法，2）解码器的大小，3）解码器的掩蔽率，4）编码器的掩蔽率。从我们的结果中我们可以得出以下观察结果。

首先，我们分析解码方法（即是否使用增强解码）的影响。可以观察到，增强解码（w.）明显优于基本解码（w.o.）。这种经验优势可以归因于增强解码的数据效率提高。在默认掩蔽率下，基本解码仅对 50% 的输入标记进行采样以进行重构，所有标记都是基于相同上下文进行预测的。使用增强解码（第 3.3 节），可以利用所有输入标记进行重构，并且每个标记都基于如式 7 所示采样的唯一上下文进行重构。这样，增强解码可以从输入数据中获得更充分、更多样化的训练信号。

其次，我们分析了解码器大小的影响，Transformer 层数从 1 增加到 3。由于增强解码仅适用于单层 Transformer，因此在这些实验中禁用了该功能。尽管计算成本更高，但扩大的解码器并没有带来任何经验收益。此外，考虑到大型解码器（#layer>1）必须禁用增强解码，这将严重损害检索性能，因此单层解码器被证明是最佳选择。

第三，我们对解码器的不同掩蔽率进行了分析，掩蔽率从 0.15 增加到 0.9。我们引入两组实验：一组使用增强解码（w.），另一组不使用增强解码（w.o.）。对于这两个实验，我们都观察到由于激进的掩蔽率导致检索质量的显著改善。对于增强解码（w.），最佳性能在 0.5 时实现；而不使用增强解码（w.o.），最佳性能在 0.7 时实现。这种微小的差异可能是因为：与可以重建所有输入标记的“w.”不同，“w.o.”仅重建被掩蔽的标记；因此，它可能以更大的掩蔽率来换取训练信号的增加。

最后，我们研究了编码器的掩蔽率。 有趣的是，与常用值 0.15 相比，略微提高掩蔽率 0.3 也能提高经验性能。对于编码器和解码器而言，重建难度的增加可能是增加掩蔽率有利于检索质量的共同原因。然而，与解码器不同，编码器的掩蔽率过高（例如 0.9）将严重损害检索性能。这是因为过大的掩蔽率会阻止生成高质量的句子嵌入，因为输入句子的大部分有用信息将被丢弃。

消融研究得出以下结论：1）通过增强解码，RetroMAE 的性能可以显著提高；2）单层 Transformer 最适合解码器；3）通过提高解码器的屏蔽率和适度提高编码器的屏蔽率，可以提高检索质量。

5 结论

我们提出了 RetroMAE，一种新颖的掩码自动编码框架，用于预训练面向检索的语言模型：输入句子被随机掩码用于编码器和解码器，句子嵌入与解码器的掩码输入相结合以重建原始输入。我们引入了不对称模型结构（全尺寸编码器和单层解码器）和不对称掩码率（编码器的比率适中，解码器的比率较高），这使得重建工作足够困难。我们还引入了增强解码，从而充分利用了预训练数据。我们在 BEIR、MS MARCO 和 Natural Question 上的实验验证了 RetroMAE 的有效性，因为与现有方法相比，它可以在零样本和监督评估方面取得显着改进。