RetroMAE 是一种基于掩码自编码器（Masked Auto-Encoder，MAE）的检索导向预训练框架，旨在提高模型在密集检索任务中的性能，以下是其详细介绍：

预训练流程

  • 输入句子两次掩码 ：输入句子在编码器和解码器端分别被不同的掩码污染。编码器根据掩码后的输入生成句子嵌入，解码器则基于句子嵌入和自身掩码后的输入，通过掩码语言建模恢复原始句子。
  • 前向传播过程 ：编码器采用多层 Transformer 结构对掩码后的输入序列进行编码，得到每个位置的隐藏表示。然后，将这些隐藏表示输入到解码器中，解码器同样是多层 Transformer 结构，它根据编码器的输出和自身掩码后的输入进行解码，输出预测的原始句子。
  • 损失函数计算 ：主要基于掩码语言建模任务，计算解码器预测结果与原始句子之间的交叉熵损失，以此来衡量模型在重建原始句子时的准确性，并通过反向传播更新模型参数。

模型结构

  • 不对称模型结构 ：编码器采用全尺寸的 BERT 风格 Transformer，能够捕捉丰富的上下文信息，生成具有较强区分性的句子嵌入。解码器则仅使用单层 Transformer，结构相对简单，使得模型在训练和推理过程中更加高效。
  • 参数细节 ：编码器通常可以是 12 层、768 个隐藏单元的 BERT 架构，解码器为 1 层 Transformer。在预训练时，编码器的输入嵌入和解码器的输入嵌入可以分别进行设计和优化。

掩码策略

  • 不对称掩码比例 ：编码器侧采用适度的掩码比例，一般为 15%-30%；解码器侧则使用激进的掩码比例，通常为 50%-70%。这种设计使解码器在重建原始句子时更具挑战性，从而迫使模型学习到更有用的信息。
  • 掩码方式选择 ：在编码器和解码器中，可以选择不同的掩码方式，如随机掩码、连续掩码等。常见的随机掩码是随机选择一定比例的 token 进行掩码，而连续掩码则是对连续的片段进行掩码。不同的掩码方式会对模型的学习效果产生一定的影响。

预训练数据

  • 数据集选择 ：RetroMAE 使用英文维基百科和 BookCorpus 作为预训练数据，这些数据集涵盖了丰富的文本内容和知识领域，能够为模型提供充分的训练信号。
  • 数据处理方式 ：在预训练数据处理过程中，会对文本进行分词、截断等操作，使其符合模型的输入要求。同时，还会对数据进行清洗和去重，以提高数据质量。

优势

  • 性能提升显著 ：在多个密集检索基准测试中表现出色，如 MS MARCO、BEIR 等，在这些数据集上的实验结果显示，RetroMAE 在监督学习和零样本学习场景下均取得了优异的性能，优于许多现有的预训练模型。
  • 预训练高效 ：由于解码器结构简单，且不需要复杂的数据增强和负采样，RetroMAE 的预训练成本与常规的 BERT / RoBERTa 风格预训练相当，具有较高的训练效率。
  • 适用性广泛 ：不仅适用于英文检索任务，还被扩展到多语言场景，如 X-RetroMAE，进一步拓展了其应用场景。

实验与应用场景

  • 实验评估 ：在 MS MARCO 数据集上，RetroMAE 的有监督密集检索性能在 MRR@10 等指标上优于其他基线方法。在 BEIR 基准测试的零样本密集检索性能上，通过 NDCG@10 指标衡量，RetroMAE 在几乎所有评估数据集上都保持了对基线方法的显著优势。
  • 实际应用 ：RetroMAE 可应用于多种检索场景，如网页搜索、问答系统、文档检索等，能够为这些任务提供更准确的文本检索和匹配能力，提升系统的性能和用户体验。