ACL 2023 | 面向信息检索的结构感知语言模型训练方法

本文中， 针对结构数据检索，我们提出了面向结构感知的稠密向量检索方法（Structure Aware DeNse ReTrievAl， SANTA），该方法设计了结构数据对齐和掩码实体预测这两个任务来继续训练预训练语言模型。实验结果表明，SANTA，它通过捕捉结构化数据的语义来学习更加精准的结构化数据表示，最终在代码和产品搜索两个任务上取得了先进的效果。

论文链接：

https://aclanthology.org/2023.findings-acl.734/

开源代码：

https://github.com/OpenMatch/SANTA

一、研究背景

结构化数据（如代码、HTML 文档和商品描述）在文章、书籍和网页中无处不在。学习文本结构背后的语义信息来表示结构化数据，对于构建更完备的检索系统至关重要。如图 1，结构化数据检索任务，如代码检索和商品检索，要求模型根据用户查询检索结构化数据。稠密向量检索是一种常用的信息检索方法，通过在向量空间中对用户查询和结构化数据进行编码，并根据向量的相似度进行匹配，来返回用户所需要的结构化数据。

然而，大多数预训练语言模型缺乏结构感知预训练，无法为结构化数据检索提供有效的向量表示。相关工作提出了一些结构感知预训练方法，用于继续训练预训练语言模型使其具备结构感知能力，来更好地表示结构化数据。这些方法通常设计了特定的掩码策略，并使用掩码语言建模对预训练语言模型进行训练。

但仅仅使用掩码语言建模可能无法充分训练预训练语言模型以进行有效的结构化数据表示。由于结构化和非结构化数据之间通常存在一些自然的对齐信号、结构化数据也蕴含有特殊的结构信息，这些为训练结构化数据表示提供了强有力的支持。我们在此基础上提出了一种结构感知语言模型预训练方法，用于实现面向结构化数据的稠密向量检索模型。

二、面向结构感知的语言模型预训练方法

针对结构数据检索，我们提出了面向结构感知的稠密向量检索方法（Structure Aware DeNse ReTrievAl， SANTA）。如图 2 所示，SANTA 设计了两个预训练任务：结构数据对齐（Structured Data Alignment，SDA）和掩码实体预测（Masked Entity Prediction，MEP）来继续训练预训练语言模型，使其对结构化数据更加敏感，更好地学习到结构化数据的表示。

1. 数据收集和处理：我们使用结构化和非结构化数据之间存在的自然对齐信号，代码-描述文档和商品描述-商品要点，来构造预训练数据对。对于代码，我们将一些代码标识符视为实体，如变量、函数名、外部库和方法，并且在 Python 和其它编程语言中分别使用 BytesIO 和 tree_sitter 来识别实体。对于商品描述，我们使用 NLTK 工具来识别同时出现在商品描述和标题中的名词及特殊名词，并将它们视为实体。
2. 结构数据对齐：我们计算编码后的非结构化数据与结构化数据之间的相似性得分，然后使用对比学习来继续训练语言模型。通过将两种模态数据对齐训练来指导语言模型优化向量空间。

1. 掩码实体预测：由于实体语义对于学习数据的结构化语义信息有着重要作用，我们在预训练语言模型时采用掩码语言模型方法来帮助语言模型捕捉数据背后的结构化语义信息。具体的，我们使用公式2训练语言模型，让其从上下文以及学习到的知识中获取必要的信息来恢复被掩码的实体，从而更好地理解数据的结构化语义信息。

三、实验结果

如表 1 所示，通过与微调后的模型性能进行比较，我们的模型（SANTA）展现出了很强的零样本能力。在代码检索任务上相比微调的 CodeT5 取得了 6.8% 的性能改进。微调之后，在代码检索和商品检索任务上分别比 CodeT5 和 T5 取得了约 8% 和 2% 的性能改进。同时，相比最先进的代码检索模型 CodeRetriever 也取得了 4.3% 的性能提升。

如表 2 所示，在基线模型上增加 MEP 任务与之前表现几乎相同，说明单纯使用掩码语言模型方法对于结构化文本的表示学习几乎没有作用。但与 MEP 不同，SDA 任务在两个结构化数据检索任务中都取得了显著的改进。当同时使用两个预训练任务后，检索性能会进一步提高。这一现象表明 MEP 任务可以通过与 SDA 任务相结合为结构化数据提供更有效的向量表示。

如图 3 所示，我们发现 SDA 任务可以很好地将非结构化数据和非结构化数据进行对齐，然而二者的向量表示会混合在一起；加入了 MEP 任务后，语言模型能够有能力对结构化和非结构文本进行区分并将其分布到不同区域中。综上，SDA 和 MEP 从不同方面帮助语言模型捕捉数据的结构化特征，进而实现更加精确的检索结果。

四、总结

当前预训练工作忽略了设计特定的结构感知预训练任务去学习结构化数据的表示，这使得他们在相应的结构数据检索任务上的表现差强人意。在这篇文章之中我们设计了结构数据对齐和掩码实体预测这两个任务，训练语言模型，使其学习数据结构背后的结构语义信息。我们的实验结果表明，SANTA，它通过捕捉结构化数据的语义来学习更加精准的结构化数据表示，最终在代码和产品搜索两个任务上取得了先进的效果。

作者：李欣泽、刘正皓等

来源：公众号【社媒派SMP】