[论文解读] LexMAE: Lexicon-Bottlenecked Pretraining for Large-Scale Retrieval

Tao Shen, Xiubo Geng+6 more|arXiv (Cornell University)|Aug 31, 2022

Multimodal Machine Learning Applications被引用 14

一句话总结

LexMAE 提出了一种新颖的预训练框架，通过引入一个词典瓶颈掩码自编码器，弥合了掩码语言建模与词典加权检索之间的差距，该模型通过连续词袋（CBoW）瓶颈学习重要性感知的词典表征。其在 BEIR 上实现了最先进（SOTA）的零样本迁移性能，并在 MS-Marco 上实现了 42.6% 的 MRR@10 和 45.8 QPS 的 CPU 推理速度，展现出极高的效率与有效性。

ABSTRACT

In large-scale retrieval, the lexicon-weighting paradigm, learning weighted sparse representations in vocabulary space, has shown promising results with high quality and low latency. Despite it deeply exploiting the lexicon-representing capability of pre-trained language models, a crucial gap remains between language modeling and lexicon-weighting retrieval -- the former preferring certain or low-entropy words whereas the latter favoring pivot or high-entropy words -- becoming the main barrier to lexicon-weighting performance for large-scale retrieval. To bridge this gap, we propose a brand-new pre-training framework, lexicon-bottlenecked masked autoencoder (LexMAE), to learn importance-aware lexicon representations. Essentially, we present a lexicon-bottlenecked module between a normal language modeling encoder and a weakened decoder, where a continuous bag-of-words bottleneck is constructed to learn a lexicon-importance distribution in an unsupervised fashion. The pre-trained LexMAE is readily transferred to the lexicon-weighting retrieval via fine-tuning. On the ad-hoc retrieval benchmark, MS-Marco, it achieves 42.6% MRR@10 with 45.8 QPS for the passage dataset and 44.4% MRR@100 with 134.8 QPS for the document dataset, by a CPU machine. And LexMAE shows state-of-the-art zero-shot transfer capability on BEIR benchmark with 12 datasets.

研究动机与目标

• 为解决掩码语言建模（MLM）与词典加权检索之间的目标不匹配问题，其中 MLM 倾向于低熵词汇，而检索则受益于高熵的枢纽词项。
• 通过在编码器与解码器之间引入连续词袋（CBoW）瓶颈，以无监督方式学习重要性感知的词典表征。
• 实现高效且低延迟的零样本与微调迁移，适用于大规模检索任务。
• 通过将词汇空间表征与以相关性为中心的语义对齐，提升预训练与下游检索目标之间的一致性。

提出的方法

• 提出一种词典瓶颈掩码自编码器（LexMAE）架构，由语言建模编码器、词典瓶颈模块和弱化解码器组成。
• 词典瓶颈模块对编码器的语言建模 logits 进行最大池化与归一化，以推导出词典重要性分布，随后将其作为连续词袋（CBoW）瓶颈中的权重使用。
• 在 CBoW 阻塞上应用基于饱和函数的 L1 范数（log(1 + ReLU)），以稳定训练并提升泛化能力。
• 使用双重目标进行模型预训练：在编码器上进行掩码语言建模，在解码器上进行弱化重建目标，采用独占掩码策略，编码器掩码比例为 30%，解码器为 50%。
• 使用词汇空间中的稀疏加权表征，对预训练模型进行端到端微调，以实现词典加权检索。
• 架构消融实验包括：禁用到词嵌入矩阵的梯度流、共享或添加独立的语言建模头，以及调整掩码策略与比例。

实验结果

研究问题

• RQ1 一种学习重要性感知词典表征的预训练框架，是否能够提升词典加权检索的性能？
• RQ2 如何缓解掩码语言建模与以相关性为导向的词典加权之间的目标不匹配问题？
• RQ3 通过 CBoW 阻塞无监督学习词典重要性分布，是否能带来更好的零样本与微调检索性能？
• RQ4 诸如阻塞类型、掩码策略与参数共享等架构选择，对 LexMAE 的检索有效性与效率有何显著影响？

主要发现

• 在 MS-Marco 文本段落检索基准上，LexMAE 实现了 42.6% 的 MRR@10，且在 CPU 上达到 45.8 QPS，展现出高效率与强大性能。
• 在 MS-Marco 文档数据集上，LexMAE 实现了 44.4% 的 MRR@100，且推理速度达 134.8 QPS，表现出优异的可扩展性与低延迟特性。
• LexMAE 在 BEIR 基准的 12 个数据集上均实现了新的最先进（SOTA）零样本迁移性能，无需微调即超越现有方法。
• 与密集 [CLS] 或无阻塞基线相比，采用基于饱和函数的 L1 范数（log(1 + ReLU)）的 CBoW 阻塞在下游检索中展现出更优的泛化能力。
• 消融实验表明，启用到词嵌入矩阵的梯度流会损害性能，而共享语言建模头或添加额外头仅产生微小影响。
• LexMAE 对多种架构与超参数选择（包括掩码策略与比例）具有鲁棒性，在多个微调阶段均保持一致的性能提升。